import sys
import csv
from pprint import pprint
from math import sqrt, log, log10, exp
import numpy as np
from scipy import optimize
from scipy.optimize import minimize
from matplotlib import pyplot as plt


#=============================
# 定数
#=============================
e  = 1.60218e-19      # C";
kB = 1.380658e-23     # JK<sup>-1</sup>";

#=============================
# 大域変数の定義
#=============================
ProgramDesc = 'LSQ fitting to diode equivalent circuit'

# CSVファイル
infile = ''

# フィッティングパラメータ初期値。
#     [I0*1e15, ndiode, Rs,     Rleak]
I0     = 3.5e-10    # diodeのprefactor
ndiode = 2.0        # 理想因子
Rs     = 1.0        # 直列抵抗
Rleak  = 1.0e7      # Shunt抵抗（リーク）


# I0, Rs, Rleakは数桁以上変わる可能性があるため、対数を取って最適化する
ai0 = [log(I0), ndiode, log(Rs), log(Rleak)]

# 同様に、フィッティングの差分関数も対数 log10(abs(I)) を取って差分を取る
# Iが 0 では対数の計算ができない。
# I = 10^-n でnが大きくなると、実際のIはほぼゼロなのに、最適化での重みが大きくなるので、
# IminよりIが小さいデータは採用しないことにする
Imin = 1.0e-15

# フィッティング範囲。Noneの場合、入力ファイルから読み込んだVの最小値、最大値になる
fitrange = [None, None]

# Callback関数で使うので、大域変数にする
Vd    = []
Id    = []
logId = []
Vi    = []
Ii    = []
logIi = []
T    = 300  # K

#=============================================
# scipy.optimize.minimizeで使うアルゴリズム
#=============================================
#nelder-mead    Downhill simplex
#powell         Modified Powell
#cg             conjugate gradient (Polak-Ribiere method)
#bfgs           BFGS法
#newton-cg      Newton-CG
#trust-ncg      信頼領域 Newton-CG 法
#dogleg         信頼領域 dog-leg 法
#L-BFGS-B’ (see here)
#TNC’ (see here)
#COBYLA’ (see here)
#SLSQP’ (see here)
#trust-constr’(see here)
#dogleg’ (see here)
#trust-exact’ (see here)
#trust-krylov’ (see here)
method = "nelder-mead"
#method = 'cg'
#method = 'powell'
#method = 'bfgs'

maxiter = 1000
tol     = 1.0e-5
h_diff  = 1.0e-3

#=============================
# グラフのフォントサイズ
#=============================
fontsize = 24
legend_fontsize = 18

#=======================================
# 起動時引数でフィッティング範囲を変更
#=======================================
argv = sys.argv
print("argv=", argv)

if len(argv) >= 2:
    infile = argv[1]
if len(argv) >= 3:
    fitrange[0] = float(argv[2])
if len(argv) >= 4:
    fitrange[1] = float(argv[3])


#=============================
# 最小化する関数の定義
#=============================

# Vが与えられたときのIを計算
# ただし、右辺に I が含まれるため、Iも与える
def Ical(ai, V, I):
#ai0 = [log(I0), ndiode, log(Rs), Log(Rleak)]
    I0 = exp(ai[0])
    if ai[1] < 0.1:
        ai[1] = 0.1
    nd = ai[1]
    Rs = exp(ai[2])
    if ai[3] > 30.0:
        ai[3] = 30.0
    Rl = exp(ai[3])

    Vd = V - Rs * I
    if V < 0.0 and Vd < V:
        Vd = V
    if V > 0.0 and Vd > V:
        Vd = V
    k = e / nd / kB / T
    return I0 * (exp(Vd * k) - 1.0) + V / Rl

# Vが与えられたときの I を計算
# optimize.fsolveを使って解となる I を求めるため、推定値 Iini を与える
def diodeI(ai, V, Iini = 0.0):
    I = optimize.fsolve(lambda I: I - Ical(ai, V, I[0]), Iini)
    return I[0]

# 各データにおける残差のリストを返す
def residual(ai, V, I):
    res = []
    Iprev = 0.0
    for i in range(len(V)):
        Ic    = diodeI(ai, V[i], Iprev)
        Iprev = Ic
        absIi = abs(I[i])
        absIc = abs(Ic)

        if absIi > Imin and absIc > Imin:
# データ、あるいは計算値の I が Imin より大きいときのみ残差を計算
            res.append(log10(absIi) - log10(absIc))
        else:
# データ、あるいは計算値の I が Imin より小さいときは残差を０にする
            res.append(0.0)

    return res

# residualで各データの残差のリストを受け取り、その二乗和を取る
def CalS2(ai):
    global Vd, Id

    res = residual(ai, Vd, Id)
    S2 = 0.0
    for i in range(len(res)):
        S2 += res[i] * res[i]
    S2 /= len(res)

    return S2

# cg法などで使う一次微分係数のリスト。有限差分により微分の近似値を求める
def diff1(ai):
    n = len(ai)
    f0 = CalS2(ai)
    df = np.empty(n)
    for i in range(n):
        aii = ai
        aii[i] = ai[i] + h_diff
        df[i] = (CalS2(aii) - f0) / h_diff
#    print("diff1: ", df)
    return df


# グラフを更新する。計算値のIc-Vcは、パラメタのリストaiから計算する
def plot(Vd, logId, Vi, logIi, ai):
    plt.clf()
    plt.title(infile, fontsize = fontsize)
    plt.xlabel(header[0], fontsize = fontsize)
    plt.ylabel('log(' + header[1] + ')', fontsize = fontsize)
    plt.plot(Vd, logId, label='raw data', marker = 'o', markersize = 0.5, linestyle = 'None')
    plt.plot(Vi, logIi, label='initial', linestyle = 'dotted')
    if ai is not None:
        Vc, Ic, logIc = CalIV(ai)
        plt.plot(Vc, logIc, label='final', linestyle = 'dashed')
    plt.legend(fontsize = legend_fontsize)
    plt.tick_params(labelsize = fontsize)
    plt.tight_layout()
    plt.pause(0.001)


# scipy.optimize.minimize() のコールバック関数
# 反復回数を記憶するため、global変数  iterを定義
iter = 0
def callback(xk):
    global iter, Vd, logId, Vi, logIi

    S2 = CalS2(xk)
# パラメータと残差二乗和を表示
    print("callback {}: xk={}  S2={}".format(iter, xk, S2))
    iter += 1
# グラフを更新
    plot(Vd, logId, Vi, logIi, xk)

# パラメータからIV特性を計算し、リストで返す
def CalIV(ai):
    Vc = []
    Ic = []
    logIc = []
    Iprev = 0.0
#    print("    %3s: %10s\t%14s" % ('i', 'V', 'I'))
    for i in range(nVc):
        V = Vcrange[0] + i * Vcstep
        I = diodeI(ai, V, Iprev)
        if abs(I) < Imin:
            pass
        else:
            Vc.append(V)
            Ic.append(I)
            logIc.append(log10(abs(I)))
        Iprev = I
    return Vc, Ic, logIc


#========================
# Main
#========================
print("")
print(ProgramDesc)

#=============================
# csvファイルの読み込み
#=============================
print("")
print("file: ", infile)

i = 0
with open(infile, "r") as f:
    reader = csv.reader(f)

    for row in reader:
        if i == 0:
            header = row
        else:
            Vi = float(row[0])
            Ii = float(row[1])
            if fitrange[0] is not None and Vi < fitrange[0]:
                continue
            if fitrange[1] is not None and fitrange[1] < Vi:
                continue
            if abs(Ii) < Imin:
                pass
            else:
                Vd.append(Vi)
                Id.append(Ii)
                logId.append(log10(abs(Ii)))
        i += 1

print("")
print("CSV data:")
print("  header:", header)
print("  V:", Vd)
print("  I:", Id)


#=============================
# 初期特性
#=============================
print("")
print("Initial parameters:")
print("  I0    =", exp(ai0[0]))
print("  ndiode=", ai0[1])
print("  Rs    =", exp(ai0[2]))
print("  Rleak =", exp(ai0[3]))

if fitrange[0] is None:
    fitrange[0] = min(Vd)
if fitrange[1] is None:
    fitrange[1] = max(Vd)

#計算データ表示範囲
Vcrange = [min(Vd), max(Vd)]
nVc = 101
Vcstep = (Vcrange[1] - Vcrange[0]) / (nVc - 1)

print("  Vc range: {} - {} at step = {}".format(Vcrange[0], Vcrange[1], Vcstep))

Vi, Ii, logIi = CalIV(ai0)

#=============================
#グラフの作成、表示
#=============================
plot(Vd, logId, Vi, logIi, None)

#=============================
# scipy.optimize()による最小化
#=============================
print("")
print("fitting range: ", fitrange)

ret = minimize(CalS2, ai0, method = method, jac = diff1, tol = tol, callback = callback,
            options = {'maxiter':maxiter, "disp":True})
print("ret=", ret)
#=======================================================================
# scipy.optimize.minimize()の戻り値はmethodにより異なるので、
# methodにより場合分けして、パラメータ ai、残差二乗和 fminを抽出
#=======================================================================
if method == 'nelder-mead':
    simplex = ret['final_simplex']
    ai   = simplex[0][0]
    fmin = ret['fun']
elif method == 'cg':
    ai = ret['x']
    fmin = ret['fun']
elif method == 'powell':
    ai = ret['x']
    fmin = ret['fun']
elif method == 'bfgs':
    ai = ret['x']
    fmin = CalS2(ai)

print("")
print("Optimized parameters: fmin=", fmin)
I0 = exp(ai[0])
nd = ai[1]
Rs = exp(ai[2])
Rl = exp(ai[3])
print("  I0    =%14.6g" % (I0))
print("  ndiode=%14.6g" % (nd))
print("  Rs    =%14.6g" % (Rs))
print("  Rleak =%14.6g" % (Rl))

#=============================
# グラフの表示
#=============================
plot(Vd, logId, Vi, logIi, ai)
print("Press ENTER to exit")
input()
