from math import exp, sqrt, log
import numpy as np
from pprint import pprint
import openpyxl


# Calculate inverse of scattering term multiplied by T^0.5
def lsqfunc(self, idx, T):
#        t = T / self.Tnorm
        t = T

        if idx == 0:
            if T == 0.0:
                return 0.0
            else:
                k = self.Eop * e / kB / T
                if k > 70.0:
                    return 0.0
                else:
                    return 1.0 / (exp(k) - 1.0)
        else:
            ip = idx - 1
            if ip >= len(self.ppolyi):
                 nfunc = int((len(self.varkeys) - 1) / 2)
                 terminate("Too many functions required (idx = {}, nfunc={})".format(idx, nfunc))

            y = pow(t, -self.ppolyi[ip])
#            print("ip=", ip, self.ppolyi[ip], self.apolyi[ip], y)

            return y

def cal_mu_components(self, T):
        if self.use_simple == 1:
            npoly = 4
        else:
            npoly = int(len(self.varkeys) / 2 + 0.00001)

        Kgb  = exp(-self.Eb * ekB / T)

        muinv_tot = 0.0
        mus  = []
        for i in range(npoly):
            if i == 0:
                muinv = self.Aop * self.lsqfunc(0, T)
                if muinv < 1.0e-10:
                    muop = 1.0e10
                else:
                    muop = 1.0 / muinv
            else:
                ip = i - 1
                muinv = self.apolyi[ip] * self.lsqfunc(i, T)
                if abs(muinv) < 1.0e-10:
                    mus.append(1.0e10)
                else:
                    mus.append(1.0 / muinv)
#                print("   mu=", 1.0 / muinv, "  muinv=", muinv, "  ip=", ip, self.apolyi[ip])

            muinv_tot += muinv
#            print("**muinv_tot=", i, muinv_tot, muinv)

        mu_ingrain = 1.0 / muinv_tot
        mu_tot = Kgb * mu_ingrain
#        print("mu_tot(0)=", mu_ingrain, muop, *mus)
        
        return mu_tot, Kgb, mu_ingrain,muop, mus

def cal_mu(self, T): #, ai = None):
#        if ai is None:
#            ai = self.ai

        mu_tot, Kgb, mu_ingrain, muop, mus = self.cal_mu_components(T)
        return mu_tot


# Simple test version only for m = 3
def mlsq_general3(self, xT, ymuinv, lsqfunc, iPrint = 1):
        print("tkMobility.mlsq_general3:")

        ndata = len(xT)
        m = 3
        print("  nData =", ndata)
        print("      xT=", strlist(xT,  "{:10.4g}"))
        print("  ymuinv=", strlist(ymuinv, "{:10.4g}"))

        Si  = np.zeros([m, 1])
        for idata in range(ndata):
            for ci in range(m):
                Si[ci, 0] += ymuinv[idata] * lsqfunc(ci, xT[idata])
    
        Sij = np.zeros([m, m])
        for ci in range(m):
             for cj in range(m):
                 for idata in range(ndata):
                     T = xT[idata]
                     Sij[ci, cj] += lsqfunc(ci, T) * lsqfunc(cj, T)

        ai = np.linalg.inv(Sij) @ Si
        ai = ai.transpose().tolist()
        ai = ai[0]
        

        """
        print("")
        print("1/mu vs lsq: lsq_general3")
        print("{:8}\t{:10}\t{:10}".format("T", "1/mu", "cal"))
        ycal = []
        for idata in range(ndata):
            T = xT[idata]
            v = 0.0
            for ci in range(m):
                v += ai[ci] * lsqfunc(ci, T)
            ycal.append(v)
            print("{:8.4g}\t{:10.4g}\t{:10.4g}".format(T, ymuinv[idata], ycal[idata]))
#        exit()
        """
        
        if iPrint == 1:
            print("tkoptimize.mslq_general3:: Vector and Matrix:") 
            print("Si=")
            pprint(Si)
            print("Sij=")
            pprint(Sij)
            print("ai=")
            pprint(ai)
            print("")

        apolyi = self.apolyi.copy()
        for i in range(3):
            apolyi[i] = ai[i]

        aalli  = [self.ai0[0]]
        ntotalvars = len(self.ai0)
        ci = 0
        for ivar in range(1, ntotalvars):
            if self.optid[ivar] == 1:
                aalli.append(ai[ci])
                ci += 1
            else:
                aalli.append(self.ai0[ivar])

        return ai, apolyi, aalli, Si, Sij

def mlsq_general(self, xT, ymuinv, optid, lsqfunc, iPrint = 1):
        print("")
        print("tkMobility.mlsq_general:")
        ndata      = len(xT)
        ntotalvars = len(self.ai0)

# 最適化するパラメータを抽出
        m = 0
        vars = []
        for i in range(1, ntotalvars):
            if self.id_llsq[i] == 1 and optid[i] == 1:
                m += 1
                vars.append(self.varkeys[i])

        print("  nTotalVars:", ntotalvars)
#        print("        Vars:", strlist(self.varkeys, "{:>10}"))
#        print("      Values:", strlist(self.ai0,     "{:10.4g}"))
        print("  id        :", strlist(optid,   "{:10d}"))
#        print("  id_llsq   :", strlist(id_llsq, "{:10d}"))
        print("    nOptVars:", m)
        print("     OptVars:", vars)
        print("  nData     :", ndata)

# 最適化する1/mu - 1/mu,fixを計算。mu,fixは、線形最小自乗に入れない他の項の和
        yconv = []
        for idata in range(ndata):
            T = xT[idata]
            ymuinvfix = 0.0
            for ivar in range(1, ntotalvars, 2):
                if optid[ivar] == 0:
                    ifunc = int((ivar - 1) / 2 + 0.00001)
#                    print("ivar=", ivar, ifunc, self.ai0[ivar])
                    ymuinvfix += self.ai0[ivar] * lsqfunc(ifunc, T)

            yconv.append(ymuinv[idata] - ymuinvfix)

#        print("            xT=", strlist(xT,  "{:10.4g}"))
#        print("        ymuinv=", strlist(ymuinv, "{:10.4g}"))
#        print("  ymuinv(conv)=", strlist(yconv, "{:10.4g}"))
        print("")
        print("{:8}\t{:12}\t{:12}\t{:12}".format("T(K)", "mu(cm2/Vs)", "1/mu(Vs/cm2)", "1/mu - 1/mu,fix"))
        for iT in range(ndata):
            print("{:8.3f}\t{:12.4g}\t{:12.4g}\t{:12.4g}".format(xT[iT], 1.0 / ymuinv[iT], ymuinv[iT], yconv[iT]))

        Si  = np.zeros([m, 1])
        for idata in range(ndata):
            T = xT[idata]
            ci = 0
            for ivar in range(1, ntotalvars):
                if optid[ivar] == 1:
                    ifunc = int((ivar - 1) / 2)
                    Si[ci, 0] += yconv[idata] * lsqfunc(ifunc, T)
                    ci += 1
#        print("Si=", Si)
    
        Sij = np.zeros([m, m])
        ci = 0
        for ivar in range(1, ntotalvars, 2):
            if optid[ivar] == 1:
                ifunc = int((ivar - 1) / 2)
                cj = 0
                for jvar in range(1, ntotalvars, 2):
                    if optid[jvar] == 1:
                        jfunc = int((jvar - 1) / 2)

                        for idata in range(ndata):
                            T = xT[idata]
                            Sij[ci, cj] += lsqfunc(ifunc, T) * lsqfunc(jfunc, T)

                        cj += 1
                ci += 1

        ai = np.linalg.inv(Sij) @ Si
        ai = ai.transpose().tolist()
        ai = ai[0]

        if iPrint == 1:
            print("")
            print("tkoptimize.mslq_general:: Vector and Matrix:") 
            print("Si=")
            pprint(Si)
            print("Sij=")
            pprint(Sij)
            print("ai=")
            pprint(ai)
            print("")

        aalli  = [self.ai0[0]]
        ci = 0
        for ivar in range(1, ntotalvars):
            if optid[ivar] == 1:
                aalli.append(ai[ci])
                ci += 1
            else:
                aalli.append(self.ai0[ivar])

        apolyi = []
        ci = 0
        for ivar in range(1, ntotalvars, 2):
            if optid[ivar] == 1:
                apolyi.append(ai[ci])
                ci += 1
            else:
                apolyi.append(self.ai0[ivar])

        if self.debug:
            print("")
            print("1/mu vs lsq: lsq_general")
            print("{:8}\t{:10}\t{:10}".format("T", "1/mu", "cal"))
            ycal = []
            for idata in range(ndata):
                T = xT[idata]
                v = 0.0
                for ci in range(m):
                    v += apolyi[ci] * lsqfunc(ci, T)
                ycal.append(v)
                print("{:8.4g}\t{:10.4g}\t{:10.4g}".format(T, ymuinv[idata], ycal[idata]))

        return ai, apolyi, aalli, Si, Sij


def main():
    print("")
    print("This is a library, not runnable")
    print("")


if __name__ == "__main__":
    main()
