"""AksharaLP MultiIndic reference. Code Apache-2.0; see NOTICE for component weight licences."""
from pathlib import Path
import json,gzip,math,re,unicodedata
import numpy as np
import torch
from bhc.xlit_q import load_any
from bhc.xlit_model import beam_search_fast, set_vocab

def norm(s):return unicodedata.normalize('NFC',s.replace('\u200c','').replace('\u200d',''))
V3=dict(w_nlm=.4,w_lm=.25,w_ch_s=1.5,w_agn=2,w_agr=0,nlm_margin=2,nlm_k=6,beam_s=8,lat_k=12,rec_n=3,use_lm=True,oov_pen_s=4.)
class MultiIndic:
 def __init__(self,root,language='hi',sentence=False,threads=1):
  self.root=Path(root);self.threads=threads;torch.set_num_threads(threads);self.cache={};self.set_language(language)
  if sentence:self.set_sentence(True)
 def set_language(self,language):
  if language not in ['hi','bn','te','kn','ta']:raise ValueError('Unsupported language')
  p=self.root/'models'/language;set_vocab(language)
  if language not in self.cache:
   if language=='hi':
    from bhc.translit2 import BhcTranslit2
    self.cache[language]=dict(hi=BhcTranslit2(str(p),threads=self.threads))
   else:
    cfg=json.loads((p/'params.json').read_text());lex={}
    for line in gzip.decompress((p/'dict.tsv.gz').read_bytes()).decode().splitlines():
     w,c=line.rsplit('\t',1);lex[w]=math.log(int(c)/cfg['dictionary_total'])
    self.cache[language]=dict(model=load_any(str(p/'xlit.q4')),p=cfg,lex=lex)
  self.language=language;self.active=self.cache[language];self.sentence=False;return self
 def set_sentence(self,enabled=True):
  if enabled and self.language!='hi':raise ValueError('Sentence mode is Hindi-only')
  if enabled and 'sent' not in self.active:
   from bhc.nlm import NLM
   from bhc.translit3 import BhcTranslit3
   p=self.root/'models/hi';nlm=NLM.load(str(p/'sentence/nlm.q8'));self.active['sent']=BhcTranslit3(str(p),threads=self.threads,nlm=nlm,**V3)
  self.sentence=enabled;return self
 @torch.inference_mode()
 def word(self,word,k=8):
  if not re.fullmatch('[A-Za-z]+',word):return [word]
  if self.language=='hi':return self.active['hi'].translit_words([word],k=k)[0]
  a=self.active;c=[];seen=set()
  for s,sc in beam_search_fast(a['model'],[word.lower()],beam=8,nbest=8)[0]:
   s=norm(s)
   if s and s not in seen:seen.add(s);c.append((s,sc))
  if not c:return []
  lex=a['lex'];r=a['p']['reranker'];mu=max((lex[s] for s,_ in c if s in lex),default=0);msc=max(sc for _,sc in c);out=[]
  for i,(s,sc) in enumerate(c):
   u=lex.get(s);inside=u is not None;f=[sc,sc/(len(s)+1),int(inside),u if inside else 0,int(i==0),i/len(c),sc-msc,u-mu if inside else 0]
   score=sum((x-m)/sd*w for x,m,sd,w in zip(f,r['mu'],r['sd'],r['w']));out.append((s,score))
  return [s for s,_ in sorted(out,key=lambda v:-v[1])[:k]]
 def transliterate(self,text):
  if self.language=='hi':
   if self.sentence:return self.active['sent'].convert(text)
   old=self.active['hi'].p['use_lm'];self.active['hi'].p['use_lm']=False
   try:return self.active['hi'].convert(text)
   finally:self.active['hi'].p['use_lm']=old
  return re.sub(r'(https?://\S+|www\.\S+|[\w.+-]+@[\w.-]+\.[A-Za-z]+|[@#][\w]+)|[A-Za-z]+',lambda m:m[0] if m[1] else (self.word(m[0],1)or[m[0]])[0],text)

 def load_word_gguf(self,path):
  from gguf_read import load_word_model
  m=load_word_model(path)
  if m.cfg.get('vocab','hi')!=self.language:raise ValueError('GGUF language does not match selected language')
  if self.language=='hi':
   self.active['hi'].model=m
   if 'sent' in self.active:self.active['sent'].model=m
  else:self.active['model']=m
  return self
 def load_sentence_gguf(self,path):
  if self.language!='hi':raise ValueError('Sentence mode is Hindi-only')
  from gguf_read import load_sentence_lm
  self.set_sentence(True);self.active['sent'].nlm=load_sentence_lm(path);return self