from stuff import *
import sys,re,argparse,pickle
from scipy.optimize import minimize
from scipy.stats import norm
from scipy.special import gammaln
from math import log,exp,sqrt,sin,pi
import numpy as np
from subprocess import Popen,PIPE
from datetime import datetime

# (Make it auto download files?)
# Get ltla.csv from https://coronavirus.data.gov.uk/api/v2/data?areaType=ltla&metric=newCasesBySpecimenDate&format=csv
# Sanger data from https://covid-surveillance-data.cog.sanger.ac.uk/download/lineages_by_ltla_and_week.tsv
# COG-UK data from https://cog-uk.s3.climb.ac.uk/phylogenetics/latest/cog_metadata.csv
# SGTF   data from Fig.16 Tech Briefing 12: https://assets.publishing.service.gov.uk/government/uploads/system/uploads/attachment_data/file/988608/Variants_of_Concern_Technical_Briefing_12_Data_England.xlsx

# LTLA_age.csv from https://api.coronavirus.data.gov.uk/v2/data?areaType=ltla&metric=newCasesBySpecimenDateAgeDemographics&format=csv

def sanitise(fn): return fn.replace(' ','_').replace("'","")

ltlaengdata=loadcsv("Local_Authority_District_to_Region__December_2020__Lookup_in_England.csv")
ltla2region=dict(zip(ltlaengdata['LAD20CD'],ltlaengdata['RGN20NM']))
ltla2name=dict(zip(ltlaengdata['LAD20CD'],map(sanitise,ltlaengdata['LAD20NM'])))
ltlapopdata=loadcsv("LTLA-NIMS-populations.csv")

# 1. Sanger uses LAD19 except E06000053 (Isles of Scilly) isn't present. I assume it's fused into E06000052 (Cornwall).
# 2. Dashboard/api use LAD19 except that E09000001 (City of London) has been fused into E09000012 (Hackney).
#    and E06000053 (Isles of Scilly) is fused into E06000052 (Cornwall).
# 3. NIMS (population and vaccine data) uses LAD20 (E070000[4567] fused into E06000060).
# Easier to fuse than to split, so standardise on LAD20 with E09000001 fused into E09000012 and E06000053 fused into E06000052.
fuseltla=dict(zip(ltlaengdata['LAD20CD'],ltlaengdata['LAD20CD']))
fuseltla['E07000004']='E06000060'
fuseltla['E07000005']='E06000060'
fuseltla['E07000006']='E06000060'
fuseltla['E07000007']='E06000060'
fuseltla['E09000001']='E09000012'
fuseltla['E06000053']='E06000052'

def coglab2uk(x): return "UK"
def coglab2country(x): return x.split('/')[0].replace('_',' ')
def coglab2coglab(x): return x
def sgtf2region(x):
  if x=='Yorkshire and Humber': return 'Yorkshire and The Humber'
  return x
def sgtf2country(x): return 'England'
def includeltla(ltla,ltlaset):
  if ltlaset=="London":
    return ltla2region[ltla]=='London'
  elif ltlaset=="test":
    return ltla2region[ltla]=='London' and ltla<'E09000010'
  elif ltlaset=="Bolton":
    return ltla=='E08000001'
  elif ltlaset=="Hartlepool":
    return ltla=='E06000001'
  elif ltlaset=="NE":
    return ltla2region[ltla]=='North East'
  elif ltlaset=="All":
    return True
  else:
    raise RuntimeError("Unrecognised ltla set "+ltlaset)

parser=argparse.ArgumentParser()
parser.add_argument('-l', '--load-options',   help='Load options from a file')
parser.add_argument('-s', '--save-options',   help='Save options to a file')
parser.add_argument('-g', '--graph-filename', help='Stem of graph filenames')
args=parser.parse_args()

# Up to 5 parameters to be optimised:
# 0: g    Difference of weekly growth rates (g(V1)-g(V0)) for unvaccinated people
# 1: rw   Weight of region counts in hierarchy
# 2: tw   Weight of total counts in hierarchy
# 3: r1   RR (= 1-VE) of vaccine for variant 1 (delta)
# 4: r0   RR (= 1-VE) of vaccine for variant 0 (alpha)

desc=["Δg_week","region_wt","eng_wt","r1","r0"]

### Model ###
### End Model ###

### Options ###

source="Sanger"
#source="COG-UK"
#source="SGTF"

# Number of parameters to optimise
N=4;r0=0.3
#N=3
#N=5

vaxeffecttime=20# Days before vaccine is presumed to have a decent effect

# Can choose location size from "LTLA", "region", "country", "UK"
# Sanger works with LTLA, region, country
# COG-UK works with country, UK
# SGTF works with region, country
locationsize="LTLA"

ltlaexclude=set()
#ltlaexclude=set(['E08000001','E12000002'])# Bolton, Manchester
ltlaset="All"
#ltlaset="London"
#ltlaset="Bolton"
#ltlaset="Hartlepool"

# Will plot graph of these locations even if only encountered during subdivision of global growth mode
specialinterest=set(['E08000001'])

mgt=5# Mean generation time in days

# Earliest day to use case data
minday=datetoday('2021-04-01')# Inclusive

# Earliest day to use VOC count data, given as end-of-week. Will be rounded up to match same day of week as lastweek.
#firstweek=minday+6
firstweek=datetoday('2021-05-01')

nif1=0.048 # Non-independence factor (1/overdispersion) for cases (less than 1 means information is downweighted)
nif2=0.255 # Non-independence factor (1/overdispersion) for VOC counts (ditto)
isd0=1.0   # Inverse sd for prior on starting number of cases of non-B.1.617.2: assume starts off similar to total number of cases
isd1=0.3   # Inverse sd for prior on starting number of cases of B.1.617.2 (0.3 is very weak)
isd2=1     # Inverse sd for prior on transmission advantage (as growth rate per day). 0 means uniform prior. 1 is very weak.

# Prior linking initial daily growth rate to estimate from pre-B.1.617.2 era
sig0=0.004

# Timescale in days over which growth rate can change significantly
# (lower = more wiggles)
bmsig=25

# Lengthscale for filtered Brownian motion
# (higher = greater amplitude for the wiggles)
bmscale=0.01

# Case ascertainment rate
asc=0.4

# Discard this many cases at the end of the list of cases by specimen day
discardcasedays=3# Pro tem to allow for Wales and NI late reporting at weekends and bank holidays

# Discard this many days of the latest COG data
discardcogdays=2

minopts={"maxiter":10000,"eps":1e-4}

mode="local growth rates"
#mode="global growth rate"
#mode="fixed growth rate",0.1

voclen=(1 if source=="COG-UK" else 7)

conf=0.95

model="scaledpoisson"
#model="NBBB"
#model="NBBB+magicprior"

### End options ###

opts={
  "Source": source,
  "Location size": locationsize,
  "LTLA set": ltlaset,
  "LTLA exclude": list(ltlaexclude),
  "Generation time (days)": mgt,
  "Earliest day for case data": daytodate(minday),
  "Earliest week (using end of week date) to use VOC count data": daytodate(firstweek),
  "Optimisation mode": mode,
  "nif1": nif1,
  "nif2": nif2,
  "Inverse sd for prior on initial non-B.1.617.2": isd0,
  "Inverse sd for prior on initial B.1.617.2": isd1,
  "Inverse sd for prior on growth": isd2,
  "Timescale of growth rate change (days)": bmsig,
  "Lengthscale for filtered Brownian motion": bmscale,
  "Case ascertainment rate": asc,
  "Number of days of case data to discard": discardcasedays,
  "Number of days of COG-UK data to discard": discardcogdays,
  "Minimiser options": minopts,
  "Length of time period over which VOC counts are given (days)": voclen,
  "Confidence level": conf,
  "Model": model
}

if args.save_options!=None:
  with open(args.save_options,'w') as fp: json.dump(opts,fp,indent=2)

if args.load_options!=None:
  with open(args.load_options,'r') as fp: lopts=json.load(fp)
  for x in lopts: opts[x]=lopts[x]

source=opts["Source"]
locationsize=opts["Location size"]
ltlaset=opts["LTLA set"]
ltlaexclude=set(opts["LTLA exclude"])
mgt=opts["Generation time (days)"]
minday=datetoday(opts["Earliest day for case data"])
firstweek=datetoday(opts["Earliest week (using end of week date) to use VOC count data"])
mode=opts["Optimisation mode"]
nif1=opts["nif1"]
nif2=opts["nif2"]
isd0=opts["Inverse sd for prior on initial non-B.1.617.2"]
isd1=opts["Inverse sd for prior on initial B.1.617.2"]
isd2=opts["Inverse sd for prior on growth"]
bmsig=opts["Timescale of growth rate change (days)"]
bmscale=opts["Lengthscale for filtered Brownian motion"]
asc=opts["Case ascertainment rate"]
discardcasedays=opts["Number of days of case data to discard"]
discardcogdays=opts["Number of days of COG-UK data to discard"]
minopts=opts["Minimiser options"]
voclen=opts["Length of time period over which VOC counts are given (days)"]
conf=opts["Confidence level"]
model=opts["Model"]

zconf=norm.ppf((1+conf)/2)

print("Options:")
print()
for x in sorted(list(opts)): print("%s:"%x,opts[x])
print()
sys.stdout.flush()

np.set_printoptions(precision=3,linewidth=150)

okplaces=set(ltla for ltla in fuseltla.values() if includeltla(ltla,ltlaset) and not ltla in ltlaexclude)
places=sorted(list(okplaces))

if source=="Sanger":
  fullsource="Wellcome Sanger Institute"
  assert voclen==7
  sanger=loadcsv("lineages_by_ltla_and_week.tsv",sep='\t')

  lastweek=datetoday(max(sanger['WeekEndDate']))
  nweeks=(lastweek-firstweek)//voclen+1
  # Sanger week number is nweeks-1-(lastweek-day)//voclen

  # Get Sanger (variant) data into a suitable form
  vocnum={}#ltla: np.zeros([nweeks,2],dtype=int) for ltla in ltla2region if ltla[0]=='E'}
  rvocnum={}
  for (date,lad19,var,n) in zip(sanger['WeekEndDate'],sanger['LTLA'],sanger['Lineage'],sanger['Count']):
    day=datetoday(date)
    week=nweeks-1-(lastweek-day)//voclen
    if week>=0 and week<nweeks:
      ltla=fuseltla[lad19]
      if ltla not in okplaces: continue
      if ltla not in vocnum: vocnum[ltla]=np.zeros([nweeks,2],dtype=int)
      if var=="B.1.617.2": vocnum[ltla][week][1]+=n
      else: vocnum[ltla][week][0]+=n
      place=ltla2region[ltla]
      if place not in rvocnum: rvocnum[place]=np.zeros([nweeks,2],dtype=int)
      if var=="B.1.617.2": rvocnum[place][week][1]+=n
      else: rvocnum[place][week][0]+=n
else: raise RuntimeError("Unrecognised source: "+source)

tvocnum=sum(rvocnum.values())

ltlapop={}
regionpop={}
for (lad20,n1,n2) in zip(ltlapopdata['LTLA Code'],ltlapopdata['Under 16'],ltlapopdata['16+']):
  ltla=fuseltla[lad20]
  if ltla not in okplaces: continue
  ltlapop[ltla]=ltlapop.get(ltla,0)+n1+n2
  region=ltla2region[ltla]
  regionpop[region]=regionpop.get(region,0)+n1+n2
rpopratio={ltla:ltlapop[ltla]/regionpop[ltla2region[ltla]] for ltla in ltlapop}
tpop=sum(regionpop.values())
tpopratio={ltla:ltlapop[ltla]/tpop for ltla in ltlapop}

# Convert daily growth rate & uncertainty into R-number-based description
# dh = 1 standard deviation
def Rdesc(h0,dh):
  (Tmin,T,Tmax)=[(exp(h*mgt)-1)*100 for h in [h0-zconf*dh,h0,h0+zconf*dh]]
  return "%.0f%% (%.0f%% - %.0f%%)"%(T,Tmin,Tmax)

# Need to scale the variables being optimised over to keep SLSQP happy
condition=np.zeros(N)+1

# Return negative log likelihood (negative because scipy can only minimise, not maximise)
# If const is true then add in all the constant terms (that don't affect the optimisation)
def NLL(xx_conditioned,const=False,pic=False):
  xx=xx_conditioned/condition
  tot=0

  # Prior on G
  #a0=log(lcases[0]+.5)
  #tot+=-((xx[0]-a0)*isd0)**2/2
  #if const: tot-=log(2*pi/isd0**2)/2

  # Prior on rweight
  #tot+=-((xx[1]-(a0-4))*isd1)**2/2
  #if const: tot-=log(2*pi/isd1**2)/2

  # Prior on tweight
  #tot+=-((xx[2]-(a0-4))*isd1)**2/2
  #if const: tot-=log(2*pi/isd1**2)/2

  # Prior on r1
  #tot+=-(xx[3]*isd2)**2/2
  #if const: tot-=log(2*pi/isd2**2)/2

  # Prior on r0
  #tot+=-(xx[4]*isd2)**2/2
  #if const: tot-=log(2*pi/isd2**2)/2

  if pic: fp=open("temp","w")
  for place in places:
    if place not in vocnum: continue
    vv=vocnum[place]
    rv=rvocnum[ltla2region[place]]
    rpr=rpopratio[place]
    tpr=tpopratio[place]
    for w in range(nweeks-1):
      rho=exp(xx[0])
      AB=vv[w]+xx[1]*rpr*rv[w]+xx[2]*tpr*tvocnum[w]
      CD=vv[w+1]
      if CD.sum()==0: continue
      if N>=4:
        r1=xx[3]
        if N>=5: r0_=xx[4]
        else: r0_=r0
        pp=pvax[place][w]
        if pp==[]:
          print(place,w,CD);raise RuntimeError("Shouldn't happen: Sanger[%s] has a case but api shows none"%place)
        #if pic: print("%s.%d  %8.5f  %12g  %12g"%(place,w,p,CD[0]/AB[0],CD[1]/AB[1]),file=fp)
      else:
        r1=r0_=0
        pp=[(1,0)]
      tl=0;tn=0
      for (n,p) in pp:
        rho_=rho*(1-p+p*r1)/(1-p+p*r0_)
        s=AB[0]+rho_*AB[1]
        #print(AB,CD,rho_,s,p,r0_,r1)
        tl+=n*(CD[0]*log(AB[0]/s)+CD[1]*log(rho_*AB[1]/s))
        tn+=n
      assert tn>0
      tot+=tl/tn
      if const: tot+=gammaln(CD[0]+CD[1]+1)-gammaln(CD[0]+1)-gammaln(CD[1]+1)# Could make a table
  if pic: fp.close()
  return -tot

def Hessian(xx):
  eps=1e-3
  H=np.zeros([N,N])
  for i in range(N-1):
    for j in range(i+1,N):
      v=0
      eps1=eps/condition[i]
      eps2=eps/condition[j]
      for (s1,s2) in [(-1,-1),(-1,1),(1,-1),(1,1)]:
        x=np.copy(xx)
        x[i]+=s1*eps1
        x[j]+=s2*eps2
        v+=s1*s2*NLL(x*condition)
      e=v/(4*eps1*eps2)
      H[i,j]=e
      H[j,i]=e
  for i in range(N):
    x=np.copy(xx)
    v=0
    eps1=eps/condition[i]
    for s in [-1,0,1]:
      x=np.copy(xx)
      x[i]+=s*eps1
      v+=(s*s*3-2)*NLL(x*condition)
    H[i,i]=v/eps1**2
  return H

# Returns log likelihood
def optimise(hint=[0.7,.25,.2,0.3,0.3][:N],statphase=False):
  xx=np.copy(hint)
  bounds=[(-5,10),(1e-2,100),(1e-2,100),(0.01,5),(0.01,5)][:N]
  res=minimize(NLL,xx*condition,bounds=bounds,method="SLSQP",options=minopts)
  if not res.success:
    print(res)
    print(place)
    print("xx =",xx)
    for x in sorted(list(opts)): print("%s:"%x,opts[x])
    print("bounds =",bounds)
    raise RuntimeError(res.message)
  xx=res.x/condition

  # Work out log likelihood including constant terms
  LL=-NLL(res.x,const=True)

  # If 'statphase', make the log likelihood a better approximation to log(integral over all parameters) using stationary phase approximation
  if statphase:
    H=Hessian(xx)
    det=np.linalg.det(H)
    if det<=0: print("Warning: Hessian not positive for %s. Can't make corrected log likelihood."%place);det=1
    LL+=N*log(2*pi)/2-log(det)/2

  # Return optimum xx log likelihood
  return res.x/condition,LL

# Assumes xx is at a local max of log likelihood
def makesamples(xx,H=None):
  if H is None: H=Hessian(xx)
  Hcond=H/condition/condition[:,None]
  eig=np.linalg.eigh(Hcond)
  # np.diag(np.matmul(np.matmul(np.transpose(eig[1]),Hcond),eig[1])) ~= eig[0]
  if not (eig[0]>0).all(): print("Hessian not +ve definite so can't do full confidence calculation");return None,None
  nsamp=100000
  t=norm.rvs(size=[nsamp,N])# nsamp x N
  sd=eig[0]**(-.5)# N
  u=t*sd# nsamp x N
  samp_cond=np.matmul(u,np.transpose(eig[1]))# nsamp x N
  samp=samp_cond/condition+xx
  cc=[]
  n0=int((1-conf)/2*nsamp)
  n1=int((1+conf)/2*nsamp)
  for i in range(N):
    a=list(samp[:,i])
    a.sort()
    cc.append((a[nsamp//2],a[n0],a[n1]))
  return samp,cc

def printplaceinfo(place,using=''):
  name=ltla2name.get(place,place)+using
  print()
  print(name)
  print("="*len(name))
  print()
  print("                        Nonvar    Var   Seen")
  for w in range(nweeks):
    day0,day1=lastweek-(nweeks-w)*voclen+1,lastweek-(nweeks-1-w)*voclen
    print(daytodate(day0),"-",daytodate(day1),"%6d %6d %6.0f"%(vocnum[place][w][0],vocnum[place][w][1],sum(cases[place][day0-minday:day1-minday+1])))
  print()

vaxdir='VaccinationData'
vaxdat={}
for x in sorted(os.listdir(vaxdir)):
  f=x.find('20')
  vaxdat[x[f:f+10]]=loadcsv(os.path.join(vaxdir,x))

def parseage_api(age):
  if '_' in age: x=age.split('_');return int(x[0]),int(x[1])+1
  if age[-1]=='+': return int(age[:-1]),150
  if age=='unassigned': return 0,150
  else: raise RuntimeError("Unrecognised age band: "+age)

# D*.0-50,D*.50-55,...,D1.0-30,D1.30-35,D1.35-40,D1.40-45,...,D2.40-45,D2.45-50,...
def parseage_nimsvax(age):
  if age[:1]!='D': return None,None
  f=age.find('.')
  g=age.find('-',f+1)
  return (age[:f],(int(age[f+1:g]),int(age[g+1:])))

# Under 16,16-29,30-34,35-39,40-44,45-49,50-54,55-59,60-64,65-69,70-74,75-79,80+,16+
def parseage_nimspop(age):
  if '-' in age: x=age.split('-');return int(x[0]),int(x[1])+1
  if age[:5]=='Under': return 0,int(age[6:])
  if age[-1]=='+': return int(age[:-1]),150
  return None

ltlaagecachedir="LTLA_age_cache"
ltlaagecachename="LTLA_age_weekly_%s_+%dweeks.pickle"%(daytodate(firstweek),nweeks)
fn=os.path.join(ltlaagecachedir,ltlaagecachename)
os.makedirs(ltlaagecachedir,exist_ok=True)
if os.path.isfile(fn):
  with open(fn,'rb') as fp:
    caseages=pickle.load(fp)
else:
  print("Loading and processing case-age data from api")
  # Todo: automate loading of LTLA_age.csv if not present or out of date.
  la=loadcsv("LTLA_age.csv")
  date0=daytodate(lastweek)
  date1=max(la['date'])
  if date1<date0: raise RuntimeError("LTLA_age.csv not up to date. Need up to %s but it ends at %s."%(date0,date1))
  caseages={}
  for lad19,date,age,cases in zip(la['areaCode'],la['date'],la['age'],la['cases']):
    ltla=fuseltla[lad19]
    day=datetoday(date)
    w=nweeks-1-(lastweek-day)//voclen
    if w>=0 and w<nweeks:
      if ltla not in caseages: caseages[ltla]={}
      a=parseage_api(age)
      if a!=(0,60) and a!=(60,150):
        caseages[ltla].setdefault(a,[0]*nweeks)[w]+=cases
  with open(fn,'wb') as fp:
    pickle.dump(caseages,fp)

from random import random,seed
#seed(42)
pvax={place:[[] for w in range(nweeks-1)] for place in places}
for w in range(nweeks-1):
  date=daytodate(firstweek+w*7+10-vaxeffecttime)
  id=max(dt for dt in vaxdat if dt<date)
  print("Using NIMS vax w/e %s to correspond to Sanger w/e %s -> w/e %s"%(id,daytodate(firstweek+w*7),daytodate(firstweek+w*7+7)))
  v=vaxdat[id]

  vaxnum={}# Map from LTLA -> age -> numvaxed
  for (i,lad20) in enumerate(v['LTLA Code']):
    ltla=fuseltla[lad20]
    if ltla in okplaces:
      if ltla not in vaxnum: vaxnum[ltla]={}
      for age in v:
        d,a=parseage_nimsvax(age)
        if a!=None:
          vaxnum[ltla][a]=vaxnum[ltla].get(a,0)+v[age][i]

  vaxpop={}# Map from LTLA -> age -> population
  for (i,lad20) in enumerate(ltlapopdata['LTLA Code']):
    ltla=fuseltla[lad20]
    if ltla in okplaces:
      if ltla not in vaxpop: vaxpop[ltla]={}
      for age in ltlapopdata:
        a=parseage_nimspop(age)
        if a!=None:
          vaxpop[ltla][a]=vaxpop[ltla].get(a,0)+ltlapopdata[age][i]

  for ltla in caseages:
    if ltla in okplaces:
      cas=caseages[ltla]
      vax=vaxnum[ltla]
      pop=vaxpop[ltla]
      pp=[]
      for ca in cas:
        n=cas[ca][w+1]
        if n==0: continue
        num=den=0
        for va in vax:
          # Want (vax number) as weighted by P(case age interval|vax age interval) = |ca intersect va|/|va|
          num+=max(min(va[1],ca[1],100)-max(va[0],ca[0]),0)/(min(va[1],100)-va[0])*vax[va]
        for pa in pop:
          den+=max(min(pa[1],ca[1],100)-max(pa[0],ca[0]),0)/(min(pa[1],100)-pa[0])*pop[pa]
        pp.append((n,min(num/den,1)))
      pvax[ltla][w]=pp

xx,L=optimise()
print()

print("Variables:",xx)
print("Log likelihood:",L)
NLL(xx*condition,const=True,pic=True)
H=Hessian(xx)
print("Hessian:");print(H)
eig=np.linalg.eigh(H)
print("Eigenvalues:",eig[0])
print()

h=xx[0]/7;dh=1/sqrt(H[0,0])/7
print("Logarithmic growth rate advantage/day: %.1f%% (%.1f%% - %.1f%%)"%(h*100,(h-zconf*dh)*100,(h+zconf*dh)*100))
print("Multiplicative growth rate advantage/day: %.1f%% (%.1f%% - %.1f%%)"%((exp(h)-1)*100,(exp(h-zconf*dh)-1)*100,(exp(h+zconf*dh)-1)*100))
print("R-number advantage: %.2f (%.2f - %.2f)"%(exp(mgt*h),exp(mgt*(h-zconf*dh)),exp(mgt*(h+zconf*dh))))
print()

print("Confidence intervals from single variables:")
for i in range(N):
  x=xx[i];dx=1/sqrt(H[i,i])
  print("%10s: %5.3f (%5.3f - %5.3f)"%(desc[i],x,x-zconf*dx,x+zconf*dx))
print()

print("Confidence intervals from multivariate calculation:")
C=np.linalg.inv(H)
for i in range(N):
  x=xx[i];dx=sqrt(C[i,i])
  print("%10s: %5.3f (%5.3f - %5.3f)"%(desc[i],x,x-zconf*dx,x+zconf*dx))
print()

samp,cc=makesamples(xx,H)
print("Confidence intervals from multivariate simulation:")
for i in range(N):
  print("%10s: %5.3f (%5.3f - %5.3f)"%((desc[i],)+cc[i]))

if 0:
  l=list(pvax)
  l.sort(key=lambda x: pvax[x][3])
  for x in l:
    print(x,"%5.0f %5.0f %5.0f %5.0f"%tuple([z*100 for z in pvax[x]]),"  ",ltla2name[x])


#!/usr/bin/env python3
#
# gw2wannier90 interface
#
# This file is distributed as part of the Wannier90 code and
# under the terms of the GNU General Public License. See the
# file `LICENSE' in the root directory of the Wannier90
# distribution, or http://www.gnu.org/copyleft/gpl.txt
#
# The webpage of the Wannier90 code is www.wannier.org
#
# The Wannier90 code is hosted on GitHub:
#
# https://github.com/wannier-developers/wannier90
#
# Designed and tested with: Quantum Espresso and Yambo
# This interface should work with any G0W0 code
# Originally written by Stepan Tsirkin
# Extended, developed and documented by Antimo Marrazzo
#
# Updated on February 19th, 2017 by Antimo Marrazzo (antimo.marrazzo@epfl.ch)
# Updated on October 7th, 2019 by Junfeng Qiao (qiaojunfeng@outlook.com)
#
import argparse
from dataclasses import dataclass
import datetime
import glob
import os
import shutil
import subprocess

import numpy as np
from scipy.io import FortranFile

def parse_args(args=None):

    parser = argparse.ArgumentParser(
        description=r"""### gw2wannier90 interface ###

Usage: gw2wannier90.py seedname options

Options can be:
  mmn, amn, spn, unk, uhu, uiu,
  spn_formatted, unk_formatted, uhu_formatted, uiu_formatted,
  write_formatted

If no options are specified, all the files are considered.

Be careful with unformatted files, they are compiler-dependent.
A safer choice is to use (bigger) formatted files, with options:
  spn_formatted, uiu_formatted, uhu_formatted, unk_formatted

In default, the output format is the same as the input format.
To generate formatted files with unformatted input, use option:
  write_formatted
""",
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument(
        "seedname",
        metavar="seedname",
        type=str,
        help="Seedname of Wannier90 files.",
    )
    parser.add_argument(
        "-o",
        "--output_seedname",
        type=str,
        help="The seedname of output files. Default is input_seedname.gw",
    )
    parser.add_argument(
        "-e",
        "--extensions",
        type=str,
        help=(
            "Comma separated list of file extensions to be converted, "
            "e.g. `-e amn,mmn` will only convert seedname.amn and seedname.mmn files. "
            "If nothing provided, all files will be converted."
        ),
    )
    parser.add_argument(
        "--no_sort",
        action="store_true",
        help="No sorting, only add GW corrections to eig.",
    )

    parsed_args = parser.parse_args(args)

    return parsed_args

def get_path_to_executable(executable: str) -> str:
    """Get path to local executable.
    :param executable: Name of executable in the $PATH variable
    :type executable: str
    :return: path to executable
    :rtype: str
    """
    path = shutil.which(executable)
    if path is None:
        raise ValueError(f"'{executable}' executable not found in PATH.")
    return path

@dataclass
class Chk:
    """Class for storing matrices in seedname.chk file."""

    header: str = None
    num_bands: int = None
    num_exclude_bands: int = None
    exclude_bands: np.ndarray = None
    real_lattice: np.ndarray = None
    recip_lattice: np.ndarray = None
    num_kpts: int = None
    mp_grid: list = None
    kpt_latt: np.ndarray = None
    nntot: int = None
    num_wann: int = None
    checkpoint: str = None
    have_disentangled: bool = None
    omega_invariant: float = None
    lwindow: np.ndarray = None
    ndimwin: np.ndarray = None
    u_matrix_opt: np.ndarray = None
    u_matrix: np.ndarray = None
    m_matrix: np.ndarray = None
    wannier_centres: np.ndarray = None
    wannier_spreads: np.ndarray = None

    def __eq__(self, other):
        if not isinstance(other, Chk):
            return NotImplemented(f"comparing {self} {other}")

        if other is self:
            return True

        eq = True

        eq = self.header == other.header
        if not eq:
            return False

        eq = self.num_bands == other.num_bands
        if not eq:
            return False

        eq = self.num_exclude_bands == other.num_exclude_bands
        if not eq:
            return False

        eq = np.allclose(self.exclude_bands, other.exclude_bands)
        if not eq:
            return False

        eq = np.allclose(self.real_lattice, other.real_lattice)
        if not eq:
            return False

        eq = np.allclose(self.recip_lattice, other.recip_lattice)
        if not eq:
            return False

        eq = self.num_kpts == other.num_kpts
        if not eq:
            return False

        eq = np.allclose(self.mp_grid, other.mp_grid)
        if not eq:
            return False

        eq = np.allclose(self.kpt_latt, other.kpt_latt)
        if not eq:
            return False

        eq = self.nntot == other.nntot
        if not eq:
            return False

        eq = self.num_wann == other.num_wann
        if not eq:
            return False

        eq = self.checkpoint == other.checkpoint
        if not eq:
            return False

        eq = self.have_disentangled == other.have_disentangled
        if not eq:
            return False

        eq = np.allclose(self.omega_invariant, other.omega_invariant)
        if not eq:
            return False

        eq = np.allclose(self.lwindow, other.lwindow)
        if not eq:
            return False

        eq = np.allclose(self.ndimwin, other.ndimwin)
        if not eq:
            return False

        eq = np.allclose(self.u_matrix_opt, other.u_matrix_opt)
        if not eq:
            return False

        eq = np.allclose(self.u_matrix, other.u_matrix)
        if not eq:
            return False

        eq = np.allclose(self.m_matrix, other.m_matrix)
        if not eq:
            return False

        eq = np.allclose(self.wannier_centres, other.wannier_centres)
        if not eq:
            return False

        eq = np.allclose(self.wannier_spreads, other.wannier_spreads)
        if not eq:
            return False

        return True

def read_chk(filename: str, formatted: bool = None, keep_temp: bool = False) -> Chk:
    """Read seedname.chk file.

    :param filename: filename
    :type filename: str
    :param formatted: defaults to None, auto detect from filename.
    :type formatted: bool, optional
    :param keep_temp: for unformatted file, creat a tempdir and run w90chk2chk.x
    in it. If True, do not remove this tempdir. Defaults to False
    :type keep_temp: bool, optional
    :return: Chk
    :rtype: Chk
    """
    import pathlib
    import tempfile

    chk = Chk()

    # From str to pathlib.Path
    filename = pathlib.Path(filename)

    if formatted is None:
        if filename.name.endswith(".chk"):
            formatted = False
        elif filename.name.endswith(".chk.fmt"):
            formatted = True
        else:
            raise ValueError(f"Cannot detect the format of {filename}")

    valid_exts = [".chk", ".chk.fmt"]
    for ext in valid_exts:
        if filename.name.endswith(ext):
            seedname = filename.name[: -len(ext)]
            break
    else:
        raise ValueError(f"{filename} not ends with {valid_exts}?")

    if not formatted:
        w90chk2chk = get_path_to_executable("w90chk2chk.x")
        tmpdir = pathlib.Path(tempfile.mkdtemp(dir="."))
        # cd tmpdir so that `w90chk2chk.log` is inside tmpdir
        os.chdir(tmpdir)
        if filename.root == "/":
            os.symlink(filename, filename.name)
        else:
            os.symlink(pathlib.Path("..") / filename, filename.name)
        call_args = [w90chk2chk, "-export", str(seedname)]
        # Some times need mpirun -n 1
        # call_args = ['mpirun', '-n', '1'] + call_args
        subprocess.check_call(call_args)
        os.chdir("..")
        filename_fmt = f"{tmpdir / filename.name}.fmt"
    else:
        filename_fmt = filename

    # Read formatted chk file
    with open(filename_fmt) as handle:
        #
        chk.header = handle.readline().strip()
        #
        chk.num_bands = int(handle.readline().strip())
        #
        chk.num_exclude_bands = int(handle.readline().strip())
        #
        chk.exclude_bands = np.zeros(chk.num_exclude_bands, dtype=int)
        #
        if chk.num_exclude_bands > 0:
            # line = handle.readline().strip().split()
            # chk.exclude_bands[:] = [int(_) for _ in line]
            for i in range(chk.num_exclude_bands):
                line = handle.readline().strip()
                chk.exclude_bands[i] = int(line)
        # Just store as a 1D array
        chk.real_lattice = np.zeros(9)
        line = handle.readline().strip().split()
        chk.real_lattice[:] = [float(_) for _ in line]
        #
        chk.recip_lattice = np.zeros(9)
        line = handle.readline().strip().split()
        chk.recip_lattice[:] = [float(_) for _ in line]
        #
        chk.num_kpts = int(handle.readline().strip())
        #
        chk.mp_grid = [int(_) for _ in handle.readline().strip().split()]
        #
        chk.kpt_latt = np.zeros((3, chk.num_kpts))
        for ik in range(chk.num_kpts):
            chk.kpt_latt[:, ik] = [float(_) for _ in handle.readline().strip().split()]
        #
        chk.nntot = int(handle.readline().strip())
        #
        chk.num_wann = int(handle.readline().strip())
        #
        chk.checkpoint = handle.readline().strip()
        # 1 -> True, 0 -> False
        chk.have_disentangled = bool(handle.readline().strip())
        if chk.have_disentangled:
            #
            chk.omega_invariant = float(handle.readline().strip())
            #
            chk.lwindow = np.zeros((chk.num_bands, chk.num_kpts), dtype=bool)
            for ik in range(chk.num_kpts):
                for ib in range(chk.num_bands):
                    # 1 -> True, 0 -> False
                    chk.lwindow[ib, ik] = bool(int(handle.readline().strip()))
            #
            chk.ndimwin = np.zeros(chk.num_kpts, dtype=int)
            for ik in range(chk.num_kpts):
                chk.ndimwin[ik] = int(handle.readline().strip())
            #
            chk.u_matrix_opt = np.zeros(
                (chk.num_bands, chk.num_wann, chk.num_kpts), dtype=complex
            )
            for ik in range(chk.num_kpts):
                for iw in range(chk.num_wann):
                    for ib in range(chk.num_bands):
                        line = [float(_) for _ in handle.readline().strip().split()]
                        chk.u_matrix_opt[ib, iw, ik] = line[0] + 1j * line[1]
        #
        chk.u_matrix = np.zeros(
            (chk.num_wann, chk.num_wann, chk.num_kpts), dtype=complex
        )
        for ik in range(chk.num_kpts):
            for iw in range(chk.num_wann):
                for ib in range(chk.num_wann):
                    line = [float(_) for _ in handle.readline().strip().split()]
                    chk.u_matrix[ib, iw, ik] = line[0] + 1j * line[1]
        #
        chk.m_matrix = np.zeros(
            (chk.num_wann, chk.num_wann, chk.nntot, chk.num_kpts), dtype=complex
        )
        for ik in range(chk.num_kpts):
            for inn in range(chk.nntot):
                for iw in range(chk.num_wann):
                    for ib in range(chk.num_wann):
                        line = [float(_) for _ in handle.readline().strip().split()]
                        chk.m_matrix[ib, iw, inn, ik] = line[0] + 1j * line[1]
        #
        chk.wannier_centres = np.zeros((3, chk.num_wann), dtype=float)
        for iw in range(chk.num_wann):
            chk.wannier_centres[:, iw] = [
                float(_) for _ in handle.readline().strip().split()
            ]
        #
        chk.wannier_spreads = np.zeros(chk.num_wann, dtype=float)
        for iw in range(chk.num_wann):
            chk.wannier_spreads[iw] = float(handle.readline().strip())

    # Read binary chk file, however its compiler dependent,
    # and it seems scipy.io.FortranFile cannot handle bool type?
    #
    #     handle = FortranFile(filename, "r")
    #     #
    #     chk.header = b"".join(handle.read_record(dtype="c"))
    #     #
    #     chk.num_bands = handle.read_record(dtype=np.int32).item()
    #     #
    #     chk.num_exclude_bands = handle.read_record(dtype=np.int32).item()
    #     #
    #     chk.exclude_bands = np.zeros(chk.num_exclude_bands, dtype=int)
    #     #
    #     if chk.num_exclude_bands > 0:
    #         line = handle.read_record(dtype=np.int32).reshape(chk.num_exclude_bands)
    #         chk.exclude_bands[:] = line[:]
    #     else:
    #         # read empty record
    #         handle.read_record(dtype=np.int32)
    #     # Just store as a 1D array
    #     chk.real_lattice = np.zeros(9)
    #     line = handle.read_record(dtype=np.float64).reshape(9)
    #     chk.real_lattice[:] = line[:]
    #     #
    #     chk.recip_lattice = np.zeros(9)
    #     line = handle.read_record(dtype=float).reshape(9)
    #     chk.recip_lattice[:] = line[:]
    #     #
    #     chk.num_kpts = handle.read_record(dtype=np.int32).item()
    #     #
    #     chk.mp_grid = handle.read_record(dtype=np.int32).reshape(3).tolist()
    #     #
    #     chk.kpt_latt = np.zeros((3, chk.num_kpts))
    #     line = handle.read_record(dtype=float).reshape((3, chk.num_kpts), order='F')
    #     chk.kpt_latt[:, :] = line[:, :]
    #     #
    #     chk.nntot = handle.read_record(dtype=np.int32).item()
    #     #
    #     chk.num_wann = handle.read_record(dtype=np.int32).item()
    #     #
    #     chk.checkpoint = b"".join(handle.read_record(dtype="c"))
    #     # 1 -> True, 0 -> False
    #     chk.have_disentangled = bool(handle.read_record(dtype=np.int32))
    #     if chk.have_disentangled:
    #         #
    #         chk.omega_invariant = handle.read_record(dtype=float).item()
    #         #
    #         chk.lwindow = np.zeros((chk.num_bands, chk.num_kpts), dtype=bool)
    #         line = handle.read_record(dtype=np.int32)
    #         line = line.reshape((chk.num_bands, chk.num_kpts), order='F')
    #         chk.lwindow[:, :] = line[:, :]
    #         #
    #         chk.ndimwin = np.array(chk.num_kpts, dtype=int)
    #         line = handle.read_record(dtype=int).reshape(chk.num_kpts)
    #         chk.ndimwin[:] = line[:]
    #         #
    #         chk.u_matrix_opt = np.array((chk.num_bands, chk.num_wann, chk.num_kpts), dtype=complex)
    #         line = handle.read_record(dtype=complex).reshape((chk.num_bands, chk.num_wann, chk.num_kpts), order='F')
    #         chk.u_matrix_opt[:, :, :] = line[:, :, :]
    #     #
    #     chk.u_matrix = np.array((chk.num_wann, chk.num_wann, chk.num_kpts), dtype=complex)
    #     line = handle.read_record(dtype=complex).reshape((chk.num_wann, chk.num_wann, chk.num_kpts), order='F')
    #     chk.u_matrix[:, :, :] = line[:, :, :]
    #     #
    #     chk.m_matrix = np.array((chk.num_wann, chk.num_wann, chk.nntot, chk.num_kpts), dtype=complex)
    #     line = handle.read_record(dtype=complex).reshape((chk.num_wann, chk.num_wann, chk.nntot, chk.num_kpts), order='F')
    #     chk.m_matrix[:, :, :, :] = line[:, :, :, :]
    #     #
    #     chk.wannier_centres = np.array((3, chk.num_wann), dtype=float)
    #     line = handle.read_record(dtype=float).reshape((3, chk.num_wann), order='F')
    #     chk.wannier_centres[:, :] = line[:, :]
    #     #
    #     chk.wannier_spreads = np.array(chk.num_wann, dtype=float)
    #     line = handle.read_record(dtype=float).reshape(chk.num_wann)
    #     chk.wannier_spreads[:] = line[:]
    #     #
    #     handle.close()

    if not formatted:
        if not keep_temp:
            shutil.rmtree(tmpdir)

    return chk

def write_chk(
    chk: Chk, filename: str, formatted: bool = None, keep_temp: bool = False
) -> None:
    """Write chk file.

    :param chk: _description_
    :type chk: Chk
    :param filename: output filename
    :type filename: str
    :param formatted: defaults to None, i.e. auto detect by filename
    :type formatted: bool, optional
    :param keep_temp: _description_, defaults to False
    :type keep_temp: bool, optional
    """
    import pathlib
    import tempfile

    # From str to pathlib.Path
    filename = pathlib.Path(filename)

    if formatted is None:
        if filename.name.endswith(".chk"):
            formatted = False
        elif filename.name.endswith(".chk.fmt"):
            formatted = True
        else:
            raise ValueError(f"Cannot detect the format of {filename}")

    valid_exts = [".chk", ".chk.fmt"]
    for ext in valid_exts:
        if filename.name.endswith(ext):
            seedname = filename.name[: -len(ext)]
            break
    else:
        raise ValueError(f"{filename} not ends with {valid_exts}?")

    if not formatted:
        tmpdir = pathlib.Path(tempfile.mkdtemp(dir="."))
        filename_fmt = f"{tmpdir / filename.name}.fmt"
    else:
        filename_fmt = filename

    # Write formatted chk file
    with open(filename_fmt, "w") as handle:
        #
        handle.write(f"{chk.header}\n")
        #
        handle.write(f"{chk.num_bands}\n")
        #
        handle.write(f"{chk.num_exclude_bands}\n")
        #
        if chk.num_exclude_bands > 0:
            # line = " ".join([str(_) for _ in chk.exclude_bands])
            # handle.write(f"{line}\n")
            for i in range(chk.num_exclude_bands):
                line = f"{chk.exclude_bands[i]}"
                handle.write(f"{line}\n")
        # Just store as a 1D array
        line = " ".join([f"{_:22.16f}" for _ in chk.real_lattice])
        handle.write(f"{line}\n")
        #
        line = " ".join([f"{_:22.16f}" for _ in chk.recip_lattice])
        handle.write(f"{line}\n")
        #
        handle.write(f"{chk.num_kpts}\n")
        #
        line = " ".join([f"{_}" for _ in chk.mp_grid])
        handle.write(f"{line}\n")
        #
        for ik in range(chk.num_kpts):
            line = " ".join([f"{_:22.16f}" for _ in chk.kpt_latt[:, ik]])
            handle.write(f"{line}\n")
        #
        handle.write(f"{chk.nntot}\n")
        #
        handle.write(f"{chk.num_wann}\n")
        #
        handle.write(f"{chk.checkpoint}\n")
        # 1 -> True, 0 -> False
        line = 1 if chk.have_disentangled else 0
        handle.write(f"{line}\n")
        if chk.have_disentangled:
            #
            handle.write(f"{chk.omega_invariant:22.16f}\n")
            #
            for ik in range(chk.num_kpts):
                for ib in range(chk.num_bands):
                    # 1 -> True, 0 -> False
                    line = 1 if chk.lwindow[ib, ik] else 0
                    handle.write(f"{line}\n")
            #
            for ik in range(chk.num_kpts):
                handle.write(f"{chk.ndimwin[ik]}\n")
            #
            for ik in range(chk.num_kpts):
                for iw in range(chk.num_wann):
                    for ib in range(chk.num_bands):
                        line = chk.u_matrix_opt[ib, iw, ik]
                        line = " ".join([f"{_:22.16f}" for _ in [line.real, line.imag]])
                        handle.write(f"{line}\n")
        #
        for ik in range(chk.num_kpts):
            for iw in range(chk.num_wann):
                for ib in range(chk.num_wann):
                    line = chk.u_matrix[ib, iw, ik]
                    line = " ".join([f"{_:22.16f}" for _ in [line.real, line.imag]])
                    handle.write(f"{line}\n")
        #
        for ik in range(chk.num_kpts):
            for inn in range(chk.nntot):
                for iw in range(chk.num_wann):
                    for ib in range(chk.num_wann):
                        line = chk.m_matrix[ib, iw, inn, ik]
                        line = " ".join([f"{_:22.16f}" for _ in [line.real, line.imag]])
                        handle.write(f"{line}\n")
        #
        for iw in range(chk.num_wann):
            line = " ".join([f"{_:22.16f}" for _ in chk.wannier_centres[:, iw]])
            handle.write(f"{line}\n")
        #
        for iw in range(chk.num_wann):
            line = f"{chk.wannier_spreads[iw]:22.16f}"
            handle.write(f"{line}\n")

    if not formatted:
        w90chk2chk = get_path_to_executable("w90chk2chk.x")
        # cd tmpdir so that `w90chk2chk.log` is inside tmpdir
        os.chdir(tmpdir)
        call_args = [w90chk2chk, "-import", str(seedname)]
        # Some times need mpirun -n 1
        # call_args = ['mpirun', '-n', '1'] + call_args
        subprocess.check_call(call_args)
        os.chdir("..")
        shutil.copy(tmpdir / f"{seedname}.chk", filename.name)

        if not keep_temp:
            shutil.rmtree(tmpdir)

def reorder_chk(seedname_in: str, seedname_out: str, bandsort: np.ndarray) -> None:
    print("----------\n CHK module  \n----------")
    filename_in = f"{seedname_in}.chk"
    filename_out = f"{seedname_out}.chk"

    if not os.path.exists(filename_in):
        print(f"WARNING: {filename_out} not written")
        return

    chk = read_chk(filename_in, formatted=False)

    # if chk.num_exclude_bands > 0:
    #     # chk.exclude_bands =
    #     # raise NotImplementedError("does not support exclude bands")

    if chk.have_disentangled:
        for ik in range(chk.num_kpts):
            chk.lwindow[:, ik] = chk.lwindow[bandsort[ik], ik]
            chk.u_matrix_opt[:, :, ik] = chk.u_matrix_opt[bandsort[ik], :, ik]
    else:
        chk.u_matrix[:, :, ik] = chk.u_matrix[bandsort[ik], :, ik]

    write_chk(chk, filename_out, formatted=False)

    print("----------\n CHK  - OK \n----------\n")

def _test_chk():
    import os
    import pathlib

    from gw2wannier90 import read_chk, write_chk

    PATH = os.environ["PATH"]
    w90_path = "/home/jqiao/git/wannier90"
    os.environ["PATH"] = f"{w90_path}:{PATH}"

    LD_LIBRARY_PATH = os.environ.get("LD_LIBRARY_PATH", "")
    mkl_path = "/opt/intel/oneapi/mpi/2021.4.0/libfabric/lib:/opt/intel/oneapi/mpi/2021.4.0/lib/release:/opt/intel/oneapi/mpi/2021.4.0/lib:/opt/intel/oneapi/mkl/2021.4.0/lib/intel64:/opt/intel/oneapi/compiler/2021.4.0/linux/lib:/opt/intel/oneapi/compiler/2021.4.0/linux/lib/x64:/opt/intel/oneapi/compiler/2021.4.0/linux/lib/emu:/opt/intel/oneapi/compiler/2021.4.0/linux/compiler/lib/intel64_lin"
    os.environ["LD_LIBRARY_PATH"] = f"{mkl_path}:{LD_LIBRARY_PATH}"

    curdir = pathlib.Path(__file__).parent

    chk = read_chk(curdir / "read_chk/silicon.chk")

    write_chk(chk, curdir / "osilicon.chk.fmt", formatted=True)

    chk2 = read_chk(curdir / "osilicon.chk.fmt")

    write_chk(chk2, curdir / "osilicon.chk")

    print(chk == chk2)

def gw2wannier90(
    seedname: str, seednameGW: str, targets: list, no_sort: bool = False
) -> None:
    print("------------------------------")
    print("##############################")
    print("### gw2wannier90 interface ###")
    print("##############################")
    print(f"Started on {datetime.datetime.now()}")

    # In case of formatted spn, uIu, uHu and UNK (mmn, amn, eig are formatted by default)
    # NB: Formatted output is strongly reccommended! Fortran binaries are compilers dependent.
    SPNformatted = "spn_formatted" in targets
    UIUformatted = "uiu_formatted" in targets
    UHUformatted = "uhu_formatted" in targets
    UNKformatted = "unk_formatted" in targets
    write_formatted = "write_formatted" in targets

    if set(targets).intersection({"spn", "uhu", "mmn", "amn", "unk", "uiu", "chk"}):
        calcAMN = "amn" in targets
        calcMMN = "mmn" in targets
        calcUHU = "uhu" in targets
        calcUIU = "uiu" in targets
        calcSPN = "spn" in targets
        calcUNK = "unk" in targets
        calcCHK = "chk" in targets
    else:
        calcAMN = True
        calcMMN = True
        calcUHU = True
        calcUIU = True
        calcSPN = True
        calcUNK = True
        calcCHK = True

    if calcUHU:
        calcMMN = True
    if calcUIU:
        calcMMN = True

    if no_sort:
        calcAMN = False
        calcMMN = False
        calcUHU = False
        calcUIU = False
        calcSPN = False
        calcUNK = False
        calcCHK = False

    # Here we open a file to dump all the intermediate steps (mainly for debugging)
    f_raw = open(seednameGW + ".gw2wannier90.raw", "w")
    # Opening seedname.nnkp file
    f = open(seedname + ".nnkp")
    # It copies the seedname.win for GW, we should make this optional
    # shutil.copy(seedname+".win",seednameGW+".win")
    while True:
        s = f.readline()
        if "begin kpoints" in s:
            break
    NKPT = int(f.readline())
    print("Kpoints number:", NKPT)
    n1 = np.array(NKPT, dtype=int)
    IKP = [
        tuple(
            np.array(
                np.round(np.array(f.readline().split(), dtype=float) * n1), dtype=int
            )
        )
        for i in range(NKPT)
    ]

    while True:
        s = f.readline()
        if "begin nnkpts" in s:
            break
    NNB = int(f.readline())

    KPNB = np.array(
        [
            [int(f.readline().split()[1]) - 1 for inb in range(NNB)]
            for ikpt in range(NKPT)
        ]
    )

    while True:
        s = f.readline()
        if "begin exclude_bands" in s:
            break
    exbands = np.array(f.readline().split(), dtype=int)
    if len(exbands) > 1 or exbands[0] != 0:
        print(
            "Exclude bands option is used: be careful to be consistent "
            "with the choice of bands for the GW QP corrections."
        )
        nexbands = exbands[0]
        exbands = np.zeros(nexbands, dtype=int)
        for i in range(nexbands):
            exbands[i] = int(f.readline().strip())
        # 0-based indexing
        exbands -= 1
    else:
        exbands = np.array([], dtype=int)

    eigenDFT = np.loadtxt(seedname + ".eig")
    nk = int(eigenDFT[:, 1].max())
    assert nk == NKPT
    nbndDFT = int(eigenDFT[:, 0].max())
    eigenDFT = eigenDFT[:, 2].reshape(NKPT, nbndDFT, order="C")
    # print(eigenDFT)
    f_raw.write("------------------------------\n")
    f_raw.write("Writing DFT eigenvalues\n")
    for line in eigenDFT:
        f_raw.write(str(line) + "\n")
    f_raw.write("------------------------------\n")

    corrections = np.loadtxt(seedname + ".gw.unsorted.eig")
    # Indexing with dict is too slow, use np.array instead.
    # corrections = {(int(l[1]) - 1, int(l[0]) - 1): l[2] for l in corrections}
    # print(corrections)
    corrections_val = np.zeros((nk, nbndDFT + len(exbands)))
    corrections_mask = np.zeros_like(corrections_val, dtype=bool)
    idx_b = corrections[:, 0].astype(int) - 1
    idx_k = corrections[:, 1].astype(int) - 1
    corrections_val[idx_k, idx_b] = corrections[:, 2]
    corrections_mask[idx_k, idx_b] = True
    # Strip excluded bands
    if len(exbands) > 0:
        corrections_val = np.delete(corrections_val, exbands, axis=1)
        corrections_mask = np.delete(corrections_mask, exbands, axis=1)
    print("G0W0 QP corrections read from ", seedname + ".gw.unsorted.eig")

    # providedGW = [
    #     ib
    #     for ib in range(nbndDFT)
    #     if all((ik, ib) in list(corrections.keys()) for ik in range(NKPT))
    # ]
    providedGW = [ib for ib in range(nbndDFT) if np.all(corrections_mask[:, ib])]
    # print(providedGW)
    f_raw.write("------------------------------\n")
    f_raw.write("List of provided GW corrections (bands indexes)\n")
    f_raw.write(str(providedGW) + "\n")
    f_raw.write("------------------------------\n")
    NBND = len(providedGW)
    print("Adding GW QP corrections to KS eigenvalues")
    # eigenDE = np.array(
    #     [[corrections[(ik, ib)] for ib in providedGW] for ik in range(NKPT)]
    # )
    # eigenDFTGW = np.array(
    #     [
    #         [eigenDFT[ik, ib] + corrections[(ik, ib)] for ib in providedGW]
    #         for ik in range(NKPT)
    #     ]
    # )
    eigenDE = corrections_val[:, providedGW]
    eigenDFTGW = eigenDFT[:, providedGW] + eigenDE

    f_raw.write("------------------------------\n")
    f_raw.write("Writing GW eigenvalues unsorted (KS + QP correction)\n")
    for line in eigenDFTGW:
        f_raw.write(str(line) + "\n")
    f_raw.write("------------------------------\n")

    if no_sort:
        print("No sorting")
    else:
        print("Sorting")
    bsort = np.array([np.argsort(eigenDFTGW[ik, :]) for ik in range(NKPT)])

    # Even if no_sort, I still output sorting list for reference
    f_raw.write("------------------------------\n")
    f_raw.write("Writing sorting list\n")
    for line in bsort:
        f_raw.write(str(line) + "\n")
    f_raw.write("------------------------------\n")

    if not no_sort:
        eigenDE = np.array([eigenDE[ik][bsort[ik]] for ik in range(NKPT)])
        eigenDFTGW = np.array([eigenDFTGW[ik][bsort[ik]] for ik in range(NKPT)])
        BANDSORT = np.array([np.array(providedGW)[bsort[ik]] for ik in range(NKPT)])

        f_raw.write("------------------------------\n")
        f_raw.write("Writing sorted GW eigenvalues\n")
        for line in eigenDFTGW:
            f_raw.write(str(line) + "\n")
        f_raw.write("------------------------------\n")

        print("GW eigenvalues sorted")

    # print eigenDFT
    print("------------------------------")
    print("writing " + seednameGW + ".eig")
    feig_out = open(seednameGW + ".eig", "w")
    for ik in range(NKPT):
        for ib in range(NBND):
            feig_out.write(f" {ib + 1:4d} {ik + 1:4d} {eigenDFTGW[ik, ib]:17.12f}\n")
    feig_out.close()
    print(seednameGW + ".eig", " written.")
    print("------------------------------\n")

    if calcAMN:
        try:
            print("----------\n AMN module  \n----------")
            f_amn_out = open(seednameGW + ".amn", "w")
            f_amn_in = open(seedname + ".amn")
            s = f_amn_in.readline().strip()
            print(s)
            f_amn_out.write(
                "{}, sorted by GW quasi-particle energies on {} \n".format(
                    s, datetime.datetime.now().isoformat()
                )
            )
            s = f_amn_in.readline()
            nb, nk, npr = np.array(s.split(), dtype=int)
            assert nk == NKPT
            assert nb == nbndDFT
            f_amn_out.write(f"  {NBND}   {nk}    {npr}   \n")

            AMN = np.loadtxt(f_amn_in, dtype=float)[:, 3:5]
            AMN = np.reshape(AMN[:, 0] + AMN[:, 1] * 1j, (nb, npr, nk), order="F")
            for ik in range(nk):
                amn = AMN[BANDSORT[ik], :, ik]
                for ipr in range(npr):
                    for ib in range(NBND):
                        f_amn_out.write(
                            " {:4d} {:4d} {:4d}  {:16.12f}  {:16.12f}\n".format(
                                ib + 1,
                                ipr + 1,
                                ik + 1,
                                amn[ib, ipr].real,
                                amn[ib, ipr].imag,
                            )
                        )
            f_amn_in.close()
            f_amn_out.close()
            print("----------\n AMN  - OK \n----------\n")
        except OSError as err:
            print(f"WARNING: {seednameGW}.amn not written : ", err)

    if calcMMN:
        try:
            print("----------\n MMN module  \n----------")

            f_mmn_out = open(os.path.join(seednameGW + ".mmn"), "w")
            f_mmn_in = open(os.path.join(seedname + ".mmn"))

            s = f_mmn_in.readline().strip()
            print(s)
            f_mmn_out.write(
                "{}, sorted by GW quasi-particle energies on {} \n".format(
                    s, datetime.datetime.now().isoformat()
                )
            )
            s = f_mmn_in.readline()
            nb, nk, nnb = np.array(s.split(), dtype=int)
            assert nb == nbndDFT
            assert nk == NKPT
            f_mmn_out.write(f"    {NBND}   {nk}    {nnb} \n")

            MMN = np.zeros((nk, nnb, NBND, NBND), dtype=complex)
            for ik in range(nk):
                for ib in range(nnb):
                    s = f_mmn_in.readline()
                    f_mmn_out.write(s)
                    ik1, ik2 = (int(i) - 1 for i in s.split()[:2])
                    assert ik == ik1
                    assert KPNB[ik][ib] == ik2
                    tmp = np.array(
                        [
                            [f_mmn_in.readline().split() for m in range(nb)]
                            for n in range(nb)
                        ],
                        dtype=str,
                    )
                    tmp = np.array(
                        tmp[BANDSORT[ik2], :, :][:, BANDSORT[ik1], :], dtype=float
                    )
                    tmp = (tmp[:, :, 0] + 1j * tmp[:, :, 1]).T
                    MMN[ik, ib, :, :] = tmp
                    for n in range(NBND):
                        for m in range(NBND):
                            f_mmn_out.write(
                                "  {:16.12f}  {:16.12f}\n".format(
                                    tmp[m, n].real, tmp[m, n].imag
                                )
                            )
            print("----------\n MMN OK  \n----------\n")
        except OSError as err:
            print(f"WARNING: {seednameGW}.mmn not written : ", err)
            if calcUHU:
                print(f"WARNING: {seednameGW}.uHu file also will not be written : ")
                calcUHU = False

    def reorder_uXu(ext, formatted=False):
        try:
            print(f"----------\n {ext} module  \n----------")

            if formatted:
                f_uXu_in = open(seedname + "." + ext)
                f_uXu_out = open(seednameGW + "." + ext, "w")
                header = f_uXu_in.readline()
                f_uXu_out.write(header)
                nbnd, NK, nnb = np.array(f_uXu_in.readline().split(), dtype=int)
                f_uXu_out.write("  ".join(str(x) for x in [NBND, NK, nnb]) + "\n")
            else:
                f_uXu_in = FortranFile(seedname + "." + ext, "r")
                header = f_uXu_in.read_record(dtype="c")
                nbnd, NK, nnb = np.array(f_uXu_in.read_record(dtype=np.int32))
                if write_formatted:
                    f_uXu_out = open(seednameGW + "." + ext, "w")
                    f_uXu_out.write("".join(header.astype(str)))
                    f_uXu_out.write("\n")
                    f_uXu_out.write("  ".join(str(x) for x in [NBND, NK, nnb]))
                    f_uXu_out.write("\n")
                else:
                    f_uXu_out = FortranFile(seednameGW + "." + ext, "w")
                    f_uXu_out.write_record(header)
                    f_uXu_out.write_record(np.array([NBND, NK, nnb], dtype=np.int32))
                header = "".join(header.astype(str))

            print(header.strip())
            print(nbnd, NK, nnb)

            assert nbnd == nbndDFT

            if formatted:
                uXu = np.loadtxt(f_uXu_in).reshape(-1)
                start = 0
                length = nbnd * nbnd

            for ik in range(NKPT):
                for ib2 in range(nnb):
                    for ib1 in range(nnb):
                        if formatted:
                            A = uXu[start : start + length]
                            start += length
                        else:
                            A = f_uXu_in.read_record(dtype=np.complex)
                        A = (
                            A.reshape(nbnd, nbnd, order="F")[
                                BANDSORT[KPNB[ik][ib2]], :
                            ][:, BANDSORT[KPNB[ik][ib1]]]
                            + np.einsum(
                                "ln,lm,l->nm",
                                MMN[ik][ib2].conj(),
                                MMN[ik][ib1],
                                eigenDE[ik],
                            )
                        ).reshape(-1, order="F")
                        if formatted or write_formatted:
                            f_uXu_out.write(
                                "".join(
                                    f"{x.real:26.16e}  {x.imag:26.16e}\n" for x in A
                                )
                            )
                        else:
                            f_uXu_out.write_record(A)
            f_uXu_out.close()
            f_uXu_in.close()
            print(f"----------\n {ext} OK  \n----------\n")
        except OSError as err:
            print(f"WARNING: {seednameGW}.{ext} not written : ", err)

    if calcUHU:
        reorder_uXu("uHu", UHUformatted)
    if calcUIU:
        reorder_uXu("uIu", UIUformatted)

    if calcSPN:
        try:
            print("----------\n SPN module  \n----------")

            if SPNformatted:
                f_spn_in = open(seedname + ".spn")
                f_spn_out = open(seednameGW + ".spn", "w")
                header = f_spn_in.readline()
                f_spn_out.write(header)
                nbnd, NK = np.array(f_spn_in.readline().split(), dtype=np.int32)
                f_spn_out.write("  ".join(str(x) for x in (NBND, NKPT)))
                f_spn_out.write("\n")
            else:
                f_spn_in = FortranFile(seedname + ".spn", "r")
                header = f_spn_in.read_record(dtype="c")
                nbnd, NK = f_spn_in.read_record(dtype=np.int32)
                if write_formatted:
                    f_spn_out = open(seednameGW + ".spn", "w")
                    f_spn_out.write("".join(header.astype(str)))
                    f_spn_out.write("\n")
                    f_spn_out.write("  ".join(str(x) for x in (NBND, NKPT)))
                    f_spn_out.write("\n")
                else:
                    f_spn_out = FortranFile(seednameGW + ".spn", "w")
                    f_spn_out.write_record(header)
                    f_spn_out.write_record(np.array([NBND, NKPT], dtype=np.int32))
                header = "".join(header.astype(str))

            print(header.strip())
            assert nbnd == nbndDFT

            indm, indn = np.tril_indices(nbnd)
            indmQP, indnQP = np.tril_indices(NBND)

            if SPNformatted:
                SPN = np.loadtxt(f_spn_in).view(complex).reshape(-1)
                start = 0
                length = (3 * nbnd * (nbnd + 1)) // 2

            for ik in range(NK):
                A = np.zeros((3, nbnd, nbnd), dtype=np.complex)
                if SPNformatted:
                    A[:, indn, indm] = SPN[start : (start + length)].reshape(
                        3, nbnd * (nbnd + 1) // 2, order="F"
                    )
                    start += length
                else:
                    A[:, indn, indm] = f_spn_in.read_record(dtype=np.complex).reshape(
                        3, nbnd * (nbnd + 1) // 2, order="F"
                    )
                A[:, indm, indn] = A[:, indn, indm].conj()
                check = np.einsum("ijj->", np.abs(A.imag))
                if check > 1e-10:
                    raise RuntimeError(f"REAL DIAG CHECK FAILED for spn: {check}")
                A = A[:, :, BANDSORT[ik]][:, BANDSORT[ik], :][
                    :, indnQP, indmQP
                ].reshape((3 * NBND * (NBND + 1) // 2), order="F")
                if SPNformatted or write_formatted:
                    f_spn_out.write(
                        "".join(f"{x.real:26.16e} {x.imag:26.16e}\n" for x in A)
                    )
                else:
                    f_spn_out.write_record(A)

            f_spn_in.close()
            f_spn_out.close()
            print("----------\n SPN OK  \n----------\n")
        except OSError as err:
            print(f"WARNING: {seednameGW}.spn not written : ", err)

    if calcUNK:
        print("----------\n UNK module  \n----------")

        unkgwdir = "UNK_GW"
        unkdftdir = "UNK_DFT"
        files_list = []
        for f_unk_name in glob.glob("UNK*.*"):
            files_list.append(f_unk_name)

        try:
            os.mkdir(unkgwdir)
            os.mkdir(unkdftdir)
        except OSError:
            pass

        for f_unk_name in files_list:
            try:
                NC = os.path.splitext(f_unk_name)[1] == ".NC"
                shutil.move("./" + f_unk_name, "./" + unkdftdir + "/")
                if UNKformatted:
                    f_unk_out = open(os.path.join(unkgwdir, f_unk_name), "w")
                    f_unk_in = open(os.path.join(unkdftdir, f_unk_name))
                    nr1, nr2, nr3, ik, nbnd = np.array(
                        f_unk_in.readline().split(), dtype=int
                    )
                    NR = nr1 * nr2 * nr3
                    if NC:
                        NR *= 2
                    f_unk_out.write(
                        " ".join(str(x) for x in (nr1, nr2, nr3, ik, NBND)) + "\n"
                    )
                    f_unk_out.write(
                        "\n".join(
                            np.array([l.rstrip() for l in f_unk_in], dtype=str)
                            .reshape((nbnd, NR), order="C")[BANDSORT[ik - 1], :]
                            .reshape(-1, order="C")
                        )
                    )
                else:
                    f_unk_in = FortranFile(os.path.join(unkdftdir, f_unk_name), "r")
                    nr1, nr2, nr3, ik, nbnd = f_unk_in.read_record(dtype=np.int32)
                    NR = nr1 * nr2 * nr3
                    unk = np.zeros((nbnd, NR), dtype=np.complex)
                    if NC:
                        unk2 = np.zeros((nbnd, NR), dtype=np.complex)
                    for ib in range(nbnd):
                        unk[ib, :] = f_unk_in.read_record(dtype=np.complex)
                        if NC:
                            unk2[ib, :] = f_unk_in.read_record(dtype=np.complex)
                    unk = unk[BANDSORT[ik - 1], :]
                    if NC:
                        unk2 = unk2[BANDSORT[ik - 1], :]
                    if write_formatted:
                        f_unk_out = open(os.path.join(unkgwdir, f_unk_name), "w")
                        f_unk_out.write(
                            " ".join(str(x) for x in (nr1, nr2, nr3, ik, NBND))
                        )
                        for i in range(NBND):
                            for j in range(NR):
                                f_unk_out.write(
                                    "\n{:21.10e} {:21.10e}".format(
                                        unk[ib, j].real, unk[ib, j].imag
                                    )
                                )
                            if NC:
                                for j in range(NR):
                                    f_unk_out.write(
                                        "\n{:21.10e} {:21.10e}".format(
                                            unk2[ib, j].real, unk2[ib, j].imag
                                        )
                                    )
                    else:
                        f_unk_out = FortranFile(os.path.join(unkgwdir, f_unk_name), "w")
                        f_unk_out.write_record(
                            np.array([nr1, nr2, nr3, ik, NBND], dtype=np.int32)
                        )
                        for i in range(NBND):
                            f_unk_out.write_record(unk[ib])
                            if NC:
                                f_unk_out.write_record(unk2[ib])
                f_unk_in.close()
                f_unk_out.close()
                shutil.move("./" + unkgwdir + "/" + f_unk_name, "./")
            except OSError as err:
                if err.errno == 21:
                    pass
                else:
                    raise err
        os.rmdir(unkgwdir)
        print(
            "UNK files have been reordered, "
            + "old files coming from DFT are available in UNK_DFT folder."
        )
        print("----------\n UNK OK  \n----------\n")

    if calcCHK:
        reorder_chk(seedname, seednameGW, BANDSORT)

    f_raw.close()

if __name__ == "__main__":
    args = parse_args()

    seedname = args.seedname  # for instance "silicon"

    if args.output_seedname is None:
        seednameGW = seedname + ".gw"  # for instance "silicon.gw"
    else:
        seednameGW = args.output_seedname

    targets = []
    if args.extensions is not None:
        targets = args.extensions.split(",")
        targets = [s.lower() for s in targets]  # options read from command line

    gw2wannier90(seedname, seednameGW, targets, args.no_sort)


#!/usr/bin/env python
# coding: utf-8

# In[1]:

from tc_python import *
import itertools as itertool
import time
import numpy as np
import matplotlib.pyplot as plt

# In[4]:

def manyPoints(database,T,P,components,phases=["bcc"],point_compositions=[(99.02,0.08)]):
    """
    Single point equilibrium calculations

    ## input:
        database name: [string],
        Temperature float,
        Pressure float,
        elements: [string],
        phases [string], if empty all phases from the database are included
        point_compositions [(float,float,...)]

    ## output: Dictionary {"stable_phases","npms","vpvs","ws","xiphs","ys", "acs","mus","bineries" }
        stable_phases: [string],
        npms: phase fractions [float],
        vpvs: volume fractions of phases [float],
        ws: weight fractions of elements [float],
        xiphs: mole fractions of elements in phases [float],
        ys: y fractions of elements is phases [float],
        bineries: binary list of all elements and stable phases [tuple (component,phase)]
        acs: activities of elements with respect to all phases
        mus: chemical potentials of all components
    """
    with TCPython() as start:
        if not phases:
            system_int = start.select_database_and_elements(database,components)
        else:
            system_int = start.select_database_and_elements(database,components).without_default_phases()
            for phase in phases:
                system_int.select_phase(phase)
        system = system_int.get_system()
        calc = system.with_single_equilibrium_calculation()

        volume_fractions,phase_fractions,weight_fractions,xs_in_phases,ys_in_phases,activities,chemical_potentials,sps,bn =         [],[],[],[],[],[],[],[],[]
        for point_composition in point_compositions:
            ticc=time.time()
            for i in range(len(components)-1):
                calc.set_condition(ThermodynamicQuantity.mole_fraction_of_a_component((components[i])), point_composition[i])
            calc.set_condition(ThermodynamicQuantity.temperature(), 1723.15)
            calc.set_condition(ThermodynamicQuantity.pressure(), 1e5)
            calc_res = calc.calculate()
            stable_phases = calc_res.get_stable_phases()
            sps.append(stable_phases)
            for phase in stable_phases:
                volume_fractions.append(calc_res.get_value_of('vpv({})'.format(phase)))
                phase_fractions.append(calc_res.get_value_of('npm({})'.format(phase)))
            for element in components:
                weight_fractions.append(calc_res.get_value_of('w({})'.format(element)))
                chemical_potentials.append(calc_res.get_value_of('mu({})'.format(element)))
            binaries = list(itertool.product(stable_phases, components))
            bn.append(binaries)
            for binary in binaries:
                xs_in_phases.append(calc_res.get_value_of('x({},{})'.format(binary[0], binary[1])))
                try:
                    ys_in_phases.append(calc_res.get_value_of('y({},{})'.format(binary[0], binary[1])))
                except Exception as error:
                    a=1
                    #ys_in_phases.append(-1)
                try:
                    activities.append(calc_res.get_value_of('ac({},{})'.format(binary[1], binary[0])))
                except Exception as error:
                    a=1
                    #ys_in_phases.append(-1)
            tocc=time.time()
            print(tocc-ticc)
        weight_fractions = np.reshape(weight_fractions,(-1,len(components)))
        chemical_potentials = np.reshape(chemical_potentials,(-1,len(components)))

        return {"stable_phases":sps,"npms":phase_fractions,                 "vpvs":volume_fractions,"ws":weight_fractions,"xiph":xs_in_phases,                 "ys":ys_in_phases,"acs":activities,"mus":chemical_potentials,"binaries":bn}

# In[5]:

help(manyPoints)

# In[6]:

database = "TCFE8"
elements = ["C","Co","N","Ti","W"]
phases = ["liquid", "fcc", "mc_shp", "graphite"]
mole_fractions=[]
for i in np.arange(0.1,0.5,0.01):
    mole_fractions.append((0.43,i,0.02,0.02,0.43))
outputs=["ws"]

# In[7]:

tic=time.time()
a=manyPoints(database,1750,1e5,elements,phases,mole_fractions)
toc = time.time()
print(toc-tic)

# In[ ]:


# coding=utf-8
from collections import defaultdict, OrderedDict
from material_parser.core.utils import simplify

import regex as re
import sympy as smp
import material_parser.core.regex_parser as rp
import material_parser.core.chemical_sets as cs
import material_parser.core.constants as cnst

from pprint import pprint

def process_formula(formula, regex_parser):
    """
    :param regex_parser:
    :param formula: str
    :return: chemical_structure
    """
    for s, r in cs.species_acronyms.items():
        formula = formula.replace(s, r)
    """
    build dictionary of elements and stoichiometries
    """
    formula_data = parse_formula(formula, regex_parser)

    """
    looking for variables in elements and stoichiometry
    """
    for el, amt in formula_data["composition"].items():
        if el not in cs.list_of_elements | formula_data["elements_x"].keys() | cnst.VACANCIES:
            formula_data["elements_x"][el] = []
        for var in re.findall(rp.re_variables, amt):
            formula_data["amounts_x"][var] = {}

    formula, \
    elements, \
    elements_x, \
    stoichiometry_x, \
    oxygen_deficiency = __refine_variables(formula_data["formula"],
                                           formula_data["composition"],
                                           formula_data["elements_x"],
                                           formula_data["amounts_x"],
                                           formula_data["oxygen_deficiency"],
                                           formula_data["oxygen_deficiency_sym"])

    if __is_acronym(formula, elements, elements_x) or __has_negative_composition(elements):
        return dict(formula=formula,
                    elements=OrderedDict(),
                    species=OrderedDict(),
                    oxygen_deficiency="",
                    phase="",
                    amounts_x={},
                    elements_x={})

    species = __get_species(formula, regex_parser) if len(elements) > 2 or formula == "H2O" else elements

    # print (formula)
    # pprint(elements)

    return dict(formula=formula,
                elements=elements,
                species=species,
                oxygen_deficiency=oxygen_deficiency,
                phase=formula_data["phase"],
                amounts_x={x: v for x, v in stoichiometry_x.items()},
                elements_x={e: v for e, v in elements_x.items()})

def parse_formula(formula, regex_parser):
    formula = formula.replace(" ", "")
    """
    separate phase, e.g. g-ABC
    """
    phase, formula = regex_parser.separate_phase(formula)

    """
    separate oxygen deficiency
    """
    formula, oxygen_deficiency, oxygen_deficiency_sym = regex_parser.separate_oxygen_deficiency(formula)

    """
    converting fractions a(b+x)/c into (a/c*b+a/c*x)
    """
    formula = regex_parser.make_fraction_convertion(formula)

    """
    check for any weird syntax (A,B)zElxEly...
    replacing with MzElxEly... and M = [A, B]
    """
    elements_x = defaultdict(str)
    stoichiometry_x = defaultdict(str)
    formula, variables = regex_parser.convert_weird_syntax(formula)
    if variables:
        elements_x["M"] = variables

    composition = __get_composition(formula)

    return dict(formula=formula,
                composition=composition,
                oxygen_deficiency=oxygen_deficiency,
                oxygen_deficiency_sym = oxygen_deficiency_sym,
                phase=phase,
                amounts_x={x: v for x, v in stoichiometry_x.items()},
                elements_x={e: v for e, v in elements_x.items()})

def __get_composition(init_formula):
    """

    :param init_formula:
    :return:
    """
    """
    if more than 4 repeating lowercase letters encountered then it is not chemical formula
    """
    if re.findall("[a-z]{4,}", init_formula):
        return OrderedDict()

    formula_dict = OrderedDict()
    formula_dict = __parse_parentheses(init_formula, "1", formula_dict)

    """
    refinement of non-variable values
    """
    incorrect = []
    for el, amt in formula_dict.items():
        formula_dict[el] = simplify(amt)
        if any(len(c) > 1 for c in re.findall("[A-Za-z]+", formula_dict[el])):
            incorrect.append(el)

    for el in incorrect:
        del formula_dict[el]

    return formula_dict

def __parse_parentheses(init_formula, init_factor, curr_dict):
    re_in_parentheses = r"\(((?>[^\(\)]+|(?R))*)\)\s*([-*\.\da-z\+/]*)"
    for m in re.finditer(re_in_parentheses, init_formula):
        factor = m.group(2) if m.group(2) != "" else "1"
        factor = simplify("(" + str(init_factor) + ")*(" + str(factor) + ")")
        unit_sym_dict = __parse_parentheses(m.group(1), factor, curr_dict)
        init_formula = init_formula.replace(m.group(0), "")

    unit_sym_dict = __get_sym_dict(init_formula, init_factor)
    for el, amt in unit_sym_dict.items():
        if el in curr_dict:
            if len(curr_dict[el]) != 0:
                curr_dict[el] = "(" + str(curr_dict[el]) + ")" + "+" + "(" + str(amt) + ")"
            else:
                curr_dict[el] = amt
        else:
            curr_dict[el] = amt

    return curr_dict

def __get_sym_dict(f, factor):
    re_sym_dict = r"([A-Z□]{1}[a-z]{0,1})\s*([\-\*\.\da-z" + "".join(cnst.GREEK_CHARS) + r"\+\/]*)"
    sym_dict = OrderedDict()

    def get_code_value(code, iterator):
        code_mapping = {"01": (iterator.group(1), iterator.group(2)),
                        "11": (iterator.group(1), iterator.group(2)),
                        "10": (iterator.group(1)[0], iterator.group(1)[1:] + iterator.group(2)),
                        "00": (iterator.group(1)[0], iterator.group(1)[1:] + iterator.group(2))}
        return code_mapping[code]

    for m in re.finditer(re_sym_dict, f):
        """
        checking for correct elements names
        """
        el_bin = "{0}{1}".format(str(int(m.group(1)[0] in cs.list_of_elements_1 | {"M"} | cnst.VACANCIES)),
                                 str(int(m.group(1) in cs.list_of_elements | {"Ln", "M"} | cnst.VACANCIES)))
        el, amt = get_code_value(el_bin, m)
        if amt.strip() == "":
            amt = "1"
        if el in sym_dict:
            sym_dict[el] = "(" + sym_dict[el] + ")" + "+" + "(" + amt + ")" + "*" + "(" + str(factor) + ")"
        else:
            sym_dict[el] = "(" + amt + ")" + "*" + "(" + str(factor) + ")"
        f = f.replace(m.group(), "", 1)
    if f.strip():
        return OrderedDict()

    """
    refinement of non-variable values
    """
    try:
        for el, amt in sym_dict.items():
            sym_dict[el] = simplify(amt)
    except:
        sym_dict = OrderedDict()
    return sym_dict

def __get_species(formula, regex_parser):
    species_in_material, species_indexs, species_dict = OrderedDict(), OrderedDict(), OrderedDict()
    material_formula = formula
    i = 0
    for species in cs.species:
        while species in material_formula:
            # print(species)
            material_formula = material_formula.replace(species, "specie" + str(i) + "_")
            species_in_material["specie" + str(i) + "_"] = species
            i += 1

    if not species_in_material:
        return OrderedDict()

    for species in cs.number_to_alphabet_dict:
        while species in material_formula:
            material_formula = material_formula.replace(species, cs.number_to_alphabet_dict[species])
            species_indexs[cs.number_to_alphabet_dict[species]] = species_in_material[species]
    species_info = parse_formula(material_formula, regex_parser)["composition"]
    for species_index in species_info:
        species_dict[species_indexs[species_index]] = species_info[species_index]
    return species_dict

def __refine_variables(formula, composition, elements_vars, stoichiometry_vars, oxy_def, oxy_def_sym):
    """
    :return:
    """
    """
    combining [RE, AE, TM] into one variable
    """
    rename_variables = [("R", "E"), ("A", "E"), ("T", "M")]
    for v1, v2 in rename_variables:
        if v1 in elements_vars and v2 in elements_vars and v1 + v2 in formula:
            elements_vars[v1 + v2] = []
            del elements_vars[v2]
            del elements_vars[v1]
            composition[v1 + v2] = composition[v2]
            del composition[v1]
            del composition[v2]

    """
    correction for Me variable
    """
    if "M" in elements_vars and "e" in stoichiometry_vars:
        elements_vars["Me"] = []
        del elements_vars["M"]
        del stoichiometry_vars["e"]
        c = composition["M"][1:]
        composition["Me"] = c if c != "" else "1.0"
        del composition["M"]

    """
    remove oxygen deficiency from variables
    """
    if not oxy_def and oxy_def_sym in stoichiometry_vars:
        oxy_def = None
    variables = [v for v in stoichiometry_vars.keys()
                 if [e for e, s in composition.items() if v in s] == ["O"]]
    oxy_def = chr(177) if len(variables) > 0 else oxy_def
    for var in variables:
        del stoichiometry_vars[var]
        composition["O"] = "1" if composition["O"] == var else composition["O"].replace(var, "").strip()
        formula = formula.replace(var, "")
    return formula, composition, elements_vars, stoichiometry_vars, oxy_def

def __is_acronym(formula, composition, variables):

    if formula in cs.ions:
        return False

    if any(ion in formula and len(ion) > 1 for ion in cs.ions):
        return False

    if len(composition) == 2 and variables:
        return True

    capital_letters = cnst.LATIN_CAPITAL - set(cs.list_of_elements_1) - {"M", "L"}
    if [r for c in capital_letters for r in re.findall(c + "[A-Z0-9\-]", formula)] \
            and all(w not in formula for w in ["RE", "OAC", "TM", "ME"]):
        return True

    if all(e.isupper() and s in ["1.0", "1"] for e, s in composition.items()):
        return True

    elements_x = [el for el in variables.keys() if len(el) == 1 and el.isupper()]
    if len(elements_x) > 1:
        return True

    if all(c.isupper() for c in formula) and any(c not in cs.list_of_elements_1 for c in formula):
        return True

    if re.findall("[A-Z]{3,}", formula) != [] and \
        all(w not in formula for w in ["CH", "COO", "OH", "NH"] + [a for a in cs.default_abbreviations.keys()]):
        return True

    if "PV" == formula[0:2]:
        return True

    return False

def __has_negative_composition(composition):

    flag = False
    try:
        flag = any(float(amt) < 0 for el, amt in composition.items())
    except:
        pass

    return flag


from typing import Union, Optional, Tuple, Dict, Any
import os

import numpy as np
import torch as th
from torch.cuda.amp import autocast

from ail.agents.rl_agent.rl_core import OnPolicyAgent
from ail.common.math import normalize
from ail.common.pytorch_util import asarray_shape2d, obs_as_tensor, disable_gradient
from ail.common.type_alias import AlgoTags, DoneMask, GymEnv, GymSpace, TensorDict

class PPO(OnPolicyAgent):
    """
    Proximal Policy Optimization algorithm (PPO) (clip version)
    Paper: https://arxiv.org/abs/1707.06347

    :param state_space: state space.
    :param action_space: action space.
    :param device: PyTorch device to which the values will be converted.
    :param seed: random seed.
    :param batch_size: size of the batch (we assume batch_size == buffer_size).
    :param policy_kwargs: arguments to be passed to the policy on creation.
        e.g. : {
            pi: [64, 64],
            vf: [64, 64],
            activation: 'relu',
            lr_actor: 3e-4,
            lr_critic: 3e-4,
            critic_type="V",
            orthogonal_init: True,
            }
    :param epoch_ppo: Number of epoch when optimizing the surrogate loss.
    :param gamma: Discount factor.
    :param clip_eps: PPO clipping parameter.
    :param coef_ent: Entropy coefficient for the loss calculation.
    :param max_grad_norm: Maximum norm for the gradient clipping.
    :param fp16: Whether to use float16 mixed precision training.
    :optim_kwargs: arguments to be passed to the optimizer.
        eg. : {
            "optim_cls": adam,
            "optim_set_to_none": True, # which set grad to None instead of zero.
            }
    :param buffer_kwargs: Arguments to be passed to the buffer.
        eg. : {
            with_reward: True,
            extra_data: ["log_pis"]
            }
    :param init_buffer: Whether to create the buffer during initialization.
    :param init_models: Whether to create the models during initialization.
    """

    def __init__(
        self,
        state_space: GymSpace,
        action_space: GymSpace,
        device: Union[th.device, str],
        seed: int,
        policy_kwargs: Dict[str, Any],
        batch_size: int = 2_000,
        epoch_ppo: int = 10,
        gamma: float = 0.99,
        gae_lambda: float = 0.97,
        clip_eps: float = 0.2,
        coef_ent: float = 0.01,
        max_grad_norm: Optional[float] = None,
        fp16: bool = False,
        optim_kwargs: Optional[dict] = None,
        buffer_kwargs: Optional[Dict[str, Any]] = None,
        init_buffer: bool = True,
        init_models: bool = True,
        expert_mode: bool = False,
        **kwargs,
    ):
        super().__init__(
            state_space,
            action_space,
            device,
            fp16,
            seed,
            gamma,
            max_grad_norm,
            batch_size,
            batch_size,  # * (Yifan) here assumes batch_size == buffer_size
            policy_kwargs,
            optim_kwargs,
            buffer_kwargs,
            init_buffer,
            init_models,
            expert_mode,
        )

        # learning rate scheduler.
        # TODO: add learning rate scheduler.
        # ? (Yifan) Is there one suitable for RL?

        """alpha_t = alpha_0 (1 - t/T)"""
        # schedule = lambda epoch: 1 - epoch/(self.param.evaluation['total_timesteps'] // self.batch_size)
        # self.scheduler_actor = optim.lr_scheduler.LambdaLR(self.optim_actor, schedule)
        # self.scheduler_critic = optim.lr_scheduler.LambdaLR(self.optim_critic, schedule)

        # Other algo params.
        self.learning_steps_ppo = 0
        self.epoch_ppo = epoch_ppo
        self.clip_eps = clip_eps
        self.gae_lambda = gae_lambda
        self.coef_ent = coef_ent

        self.tag = AlgoTags.PPO

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}"

    def is_update(self, step: int) -> bool:
        """Whether or not to update the agent"""
        return step % self.batch_size == 0

    def step(
        self,
        env: GymEnv,
        state: th.Tensor,
        episode_timesteps: th.Tensor,
        global_timesteps: Optional[int] = None,
        add_absorbing_state: bool = False,
    ) -> Tuple[np.ndarray, int]:
        """
        Intereact with environment and store the transition.

        :param env: gym environment
        :param state: orginal state return by the environment
        :param episode_timesteps: number of timesteps this episode
        :param total_timesteps: total number of timesteps to run in outer loop
        :return: next_state, episode length
        """
        episode_timesteps += 1

        # Sample actions from action distribution.
        # which is then wrapped by tanh transform to keep it in range [-1, 1].
        action, log_pi = self.explore(obs_as_tensor(state, self.device), scale=False)

        # Resacle actions to match original action space.
        scale_action = (
            self.scale_action(action) if not self.normalized_action_space else action
        )

        # Interact with environment (Info might be useful for some special env).
        next_state, reward, done, info = env.step(scale_action)

        # Done mask removes the time limit constrain of some env to keep makorvian.
        # Agent keeps alive should not be done by env's time limit.
        # See: https://github.com/sfujim/TD3/blob/master/main.py#L127
        # * Here we use an inverse convention in which DONE = 0 and NOT_DONE = 1
        # * to match absorbing state implementation in DAC paper.
        done_mask: float
        if (episode_timesteps == env._max_episode_steps) or not done:
            done_mask = DoneMask.NOT_DONE.value
        else:
            done_mask = DoneMask.DONE.value

        absorbing_cond = all(
            [add_absorbing_state, done, episode_timesteps < env._max_episode_steps]
        )
        if absorbing_cond:
            next_state = env.get_absorbing_state()

        data = {
            "obs": asarray_shape2d(state),
            "acts": asarray_shape2d(action),
            "rews": asarray_shape2d(reward),
            "dones": asarray_shape2d(done_mask),
            "log_pis": asarray_shape2d(log_pi),
            "next_obs": asarray_shape2d(next_state),
        }

        # Store transition.
        # * NOT ALLOW size larger than buffer capcity.
        self.buffer.store(data, truncate_ok=False)

        # Reset env if encounter done signal (not done mask!)
        if done:
            episode_timesteps = 0
            next_state = env.reset()
            # Add a absorbing state to buffer when done.
            if add_absorbing_state and (episode_timesteps < env._max_episode_steps):
                # A fake action for the absorbing state.
                zero_action = np.zeros(env.action_space.shape)
                absorbing_state = env.get_absorbing_state()
                absorbing_data = {
                    "obs": asarray_shape2d(absorbing_state),
                    "acts": asarray_shape2d(zero_action),
                    "rews": asarray_shape2d(0.0),
                    "dones": asarray_shape2d(DoneMask.ABSORBING.value),
                    "log_pis": asarray_shape2d(log_pi),  # TODO: what to do with log_pi?
                    "next_obs": asarray_shape2d(absorbing_state),
                }
                self.buffer.store(absorbing_data, truncate_ok=False)
        return next_state, episode_timesteps

    def update(self, log_this_batch: bool = False) -> Dict[str, Any]:
        """
        A general road map for updating the model.
        Obtain the training batch and perform update.
        :return train_logs: dict of training logs
        """
        self.learning_steps += 1
        rollout_data = self.buffer.get()

        # Clear buffer after getting entire buffer.
        self.buffer.reset()
        train_logs = self.update_algo(rollout_data, log_this_batch)
        return train_logs

    def update_algo(
        self, data: TensorDict, log_this_batch: bool = False
    ) -> Dict[str, Any]:
        """
        Update the actor and critic.
        :param data: a batch of randomly sampled transitions
        :return train_logs: dict of training logs
        """
        states, actions, rewards, dones, next_states, log_pis = (
            data["obs"],
            data["acts"],
            data["rews"],
            data["dones"],
            data["next_obs"],
            data["log_pis"],
        )
        with th.no_grad():
            values = self.critic(states)
            next_values = self.critic(next_states)

        targets, gaes = calculate_gae(
            rewards, (1.0 - dones), values, next_values, self.gamma, self.gae_lambda
        )

        for _ in range(self.epoch_ppo):
            self.learning_steps_ppo += 1
            loss_critic = self._update_critic(states, targets)
            loss_actor, pi_info = self._update_actor(states, actions, log_pis, gaes)

        if log_this_batch:
            # Return log changes(key used for logging name).
            return {
                "actor_loss": loss_actor,
                "critic_loss": loss_critic,
                "approx_kl": pi_info["kl"],
                "entropy": pi_info["ent"],
                "clip_fraction": pi_info["cf"],
                "pi_lr": self.lr_actor,
                "vf_lr": self.lr_critic,
                "learn_steps_ppo": self.learning_steps_ppo,
            }
        else:
            return {}

    def _update_critic(self, states: th.Tensor, targets: th.Tensor) -> th.Tensor:
        """
        Update critic. (value function approximation)
        :param states:
        :param targets: should be gae + v_pred
        return: critic loss
        """
        self.optim_critic.zero_grad(set_to_none=self.optim_set_to_none)
        with autocast(enabled=self.fp16):
            loss_critic = (self.critic(states) - targets).pow(2).mean()
        self.one_gradient_step(loss_critic, self.optim_critic, self.critic)
        return loss_critic.detach()

    def _update_actor(
        self,
        states: th.Tensor,
        actions: th.Tensor,
        log_pis_old: th.Tensor,
        gaes: th.Tensor,
    ) -> Tuple[th.Tensor, Dict[str, Any]]:
        """
        Update actor. (function for computing PPO policy loss)
        :param states:
        :param actions:
        :param log_pis_old:
        :param gaes: general advantage estimation
        : return: actor loss, policy_info
        """
        log_pis = self.actor.evaluate_log_pi(states, actions)

        # * (Yifan) Since we bounded the mean action with tanh(),
        # * there is no analytical form of entropy
        # Approximate entropy.
        approx_ent = -log_pis.mean()

        # ratio between old and new policy, should be one at the first iteration
        log_ratios = log_pis - log_pis_old
        ratios = (log_ratios).exp()

        # clipped surrogate loss
        loss_actor1 = ratios * gaes
        loss_actor2 = th.clamp(ratios, 1.0 - self.clip_eps, 1.0 + self.clip_eps) * gaes
        loss_actor = -th.min(loss_actor1, loss_actor2).mean()

        self.optim_actor.zero_grad(set_to_none=self.optim_set_to_none)
        with autocast(enabled=self.fp16):
            loss_actor_ent = loss_actor - self.coef_ent * approx_ent
        self.one_gradient_step(loss_actor_ent, self.optim_actor, self.actor)

        """
        Calculate approximate form of reverse KL Divergence for early stopping.
        See issue #417: https://github.com/DLR-RM/stable-baselines3/issues/417
        and discussion in PR #419: https://github.com/DLR-RM/stable-baselines3/pull/419
        and Schulman blog: https://joschu.net/blog/kl-approx.html
        KL(q||p): (r-1) - log(r), where r = p(x)/q(x)
        """
        # ! (Yifan) Deprecated :
        # ! Naive version: approx_kl = (log_pi_old - log_pi).mean().item()
        # ! This is an unbiased estimator, but it has large variance.
        # ! Since it can take on negative values.
        # ! as opposed to the actual KL Divergence measure
        # Useful extra info
        with th.no_grad():
            approx_kl = ((ratios - 1) - log_ratios).mean()
            clipped = ratios.gt(1 + self.clip_eps) | ratios.lt(1 - self.clip_eps)
            clip_frac = th.as_tensor(clipped, dtype=th.float32).mean()
            pi_info = {"kl": approx_kl, "ent": approx_ent.detach(), "cf": clip_frac}
        return loss_actor.detach(), pi_info

    def save_models(self, save_dir: str) -> None:
        """
        Save the model. (Only save actor to reduce workloads)
        """
        super().save_models(save_dir)
        th.save(self.actor.state_dict(), os.path.join(save_dir, "actor.pth"))

    @classmethod
    def load(
        cls,
        path: str,
        policy_kwargs: Dict[str, Any],
        env: Union[GymEnv, str, None] = None,
        state_space: Optional[GymSpace] = None,
        action_space: Optional[GymSpace] = None,
        device: Union[th.device, str] = "cpu",
        seed: int = 42,
        **kwargs,
    ) -> "PPO":
        """
        Load the model from a saved model directory.
        we only load actor.
        """
        super().load(env, state_space, action_space)
        if env is not None:
            if isinstance(env, str):
                import gym

                env = gym.make(env)
            state_space, action_space = env.observation_space, env.action_space

        ppo_expert = cls(
            state_space,
            action_space,
            device,
            seed,
            policy_kwargs,
            init_buffer=False,
            init_models=False,
            expert_mode=True,
            **kwargs,
        )
        state_dict = th.load(path)
        ppo_expert.actor.load_state_dict(state_dict)
        disable_gradient(ppo_expert.actor)
        ppo_expert.actor.eval()
        return ppo_expert

def calculate_gae(
    rewards: th.Tensor,
    dones: th.Tensor,
    values: th.Tensor,
    next_values: th.Tensor,
    gamma: float,
    lambd: float,
    normal: bool = True,
) -> Tuple[th.Tensor, th.Tensor]:
    """
    Compute the lambda-return (TD(lambda) estimate) and GAE(lambda) advantage.

    Uses Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438)
    to compute the advantage. To obtain vanilla advantage (A(s) = R - V(S))
    where R is the discounted reward with value bootstrap,
    set `lambd=1.0`.
    https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/ppo1/pposgd_simple.py#L66
    """
    # Calculate TD errors.
    deltas = rewards + gamma * next_values * (1.0 - dones) - values
    # Initialize gae.
    gaes = th.empty_like(rewards)

    # Calculate gae recursively from behind.
    gaes[-1] = deltas[-1]
    for t in reversed(range(rewards.size(0) - 1)):
        gaes[t] = deltas[t] + gamma * lambd * (1.0 - dones[t]) * gaes[t + 1]

    targets = values + gaes

    if normal:
        return targets, normalize(gaes, gaes.mean(), gaes.std())
    else:
        return targets, gaes


import os
import sys
import numpy as np

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule

class gpu_context:
    """
    GPU context which manages the allocation of memory, the movement of memory between python and the GPU,
    and the calling of GPU funcitons

    Written: Devin Cody, 2021
    """
    def __init__(self, len_gpu_arrays = 10000000):
        self.gpu_initalized = False
        self.len_gpu_arrays = len_gpu_arrays

        try:
            print("Compiling kernel")
            if "win" in sys.platform:
                f_newton = open(os.path.dirname(__file__) + "\\kernels\\newton.cu", 'r')
                f_mikkola = open(os.path.dirname(__file__) + "\\kernels\\mikkola.cu", 'r')
            else:
                f_newton = open(os.path.dirname(__file__) + "/kernels/newton.cu", 'r')
                f_mikkola = open(os.path.dirname(__file__) + "/kernels/mikkola.cu", 'r')

            fstr_newton = "".join(f_newton.readlines())
            mod_newton = SourceModule(fstr_newton)
            self.newton_gpu = mod_newton.get_function("newton_gpu")

            fstr_mikkola = "".join(f_mikkola.readlines())
            mod_mikkola = SourceModule(fstr_mikkola)
            self.mikkola_gpu = mod_mikkola.get_function("mikkola_gpu")

            print("Allocating with {} bytes".format(self.len_gpu_arrays))
            self.tolerance = np.array([1e-9], dtype = np.float64)
            self.max_iter = np.array([100])
            self.eanom = None

            self.d_manom = cuda.mem_alloc(self.len_gpu_arrays)
            self.d_ecc = cuda.mem_alloc(self.len_gpu_arrays)
            self.d_eanom = cuda.mem_alloc(self.len_gpu_arrays)

            self.d_tol = cuda.mem_alloc(self.tolerance.nbytes)
            self.d_max_iter = cuda.mem_alloc(self.max_iter.nbytes)

            print("Copying parameters to GPU")
            cuda.memcpy_htod(self.d_tol, self.tolerance)
            cuda.memcpy_htod(self.d_max_iter, self.max_iter)
            gpu_initalized = True
        except Exception as e:
            print("Error: KEPLER: Unable to initialize Kepler GPU solver context")
            raise(e)

    def newton(self, manom, ecc, eanom, eanom0 = None, tolerance=1e-9, max_iter=100):
        """
        Moves numpy arrays onto the GPU memory, calls the Newton-Raphson solver for eccentric anomaly
        and copies the result back into a numpy array.

        Args:
            manom (np.array): array of mean anomalies
            ecc (np.array): array of eccentricities
            eanom (np.array): array of eccentric anomalies (return by reference)
            eanom0 (np.array, optional): array of first guess for eccentric anomaly, same shape as manom (optional)
        Return:
            None: eanom is changed by reference

        Written: Devin Cody, 2021

        """
        # Check to make sure we have enough data to process orbits
        if (self.len_gpu_arrays < manom.nbytes):
            self.len_gpu_arrays = manom.nbytes
            self.d_manom = cuda.mem_alloc(self.len_gpu_arrays)
            self.d_ecc = cuda.mem_alloc(self.len_gpu_arrays)
            self.d_eanom = cuda.mem_alloc(self.len_gpu_arrays)

        cuda.memcpy_htod(self.d_manom, manom)
        cuda.memcpy_htod(self.d_ecc, ecc)
        cuda.memcpy_htod(self.d_tol, tolerance)
        cuda.memcpy_htod(self.d_max_iter, max_iter)

        # Initialize at E=M, E=pi is better at very high eccentricities
        if eanom0 is None:
            cuda.memcpy_dtod(self.d_eanom, self.d_manom, self.len_gpu_arrays)
        else:
            cuda.memcpy_htod(self.d_eanom, eanom0)

        self.newton_gpu(self.d_manom, self.d_ecc, self.d_eanom, self.d_max_iter, self.d_tol, grid = (len(manom)//64+1,1,1), block = (64,1,1))
        cuda.memcpy_dtoh(eanom, self.d_eanom)

    def mikkola(self, manom, ecc, eanom):
        """
        Moves numpy arrays onto the GPU memory, calls the analtyical Mikkola solver for eccentric anomaly
        and copies the result back into a numpy array.

        Args:
            manom (np.array): array of mean anomalies between 0 and 2pi
            ecc (np.array): eccentricity
            eanom (np.array): array of eccentric anomalies (return by reference)
        Return:
            None: eanom is changed by reference

        Written: Devin Cody, 2021
        """
        # Check to make sure we have enough data to process orbits
        if (self.len_gpu_arrays < manom.nbytes):
            self.len_gpu_arrays = manom.nbytes
            self.d_manom = cuda.mem_alloc(self.len_gpu_arrays)
            self.d_ecc = cuda.mem_alloc(self.len_gpu_arrays)

        cuda.memcpy_htod(self.d_manom, manom)
        cuda.memcpy_htod(self.d_ecc, ecc)

        self.mikkola_gpu(self.d_manom, self.d_ecc, self.d_eanom, grid = (len(manom)//64+1,1,1), block = (64,1,1))
        cuda.memcpy_dtoh(eanom, self.d_eanom)