#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Estimate saturation mobility from TFT data by using first-order polynomial least-squares fitting
(with error estimation using degree of freedom n-2).

This refactored code reads data from a CSV file, performs linear regression on sqrt(Ids) vs. Vgs,
and plots various figures including the transfer and output characteristics, error bars,
and the dependency of extracted parameters on the fitting range.
"""

import sys
import re
import numpy as np
import csv
from math import sqrt, pi
from matplotlib import pyplot as plt


#===================================
# physical constants
#===================================
PI    = pi
PI2   = 2.0 * pi
h     = 6.6260755e-34    # Planck's constant [Js]
hbar  = 1.05459e-34      # Reduced Planck's constant [Js]
c     = 2.99792458e8     # Speed of light [m/s]
e     = 1.60218e-19      # Elementary charge [C]
e0    = 8.854418782e-12  # Vacuum permittivity [F/m]
kB    = 1.380658e-23     # Boltzmann constant [J/K]
me    = 9.1093897e-31    # Electron mass [kg]
R     = 8.314462618      # Gas constant [J/K/mol]
a0    = 5.29177e-11      # Bohr radius [m]


#===================================
# parameters
#===================================
infile = 'TransferCurve.csv'  # CSV input file

dg   = 100.0e-9  # m, gate oxide thickness
erg  = 11.9      # Relative permittivity

W    = 300.0e-6  # m, channel width
L    =  50.0e-6  # m, channel length

Vds0   = 10.0   # V, Vds value to be used for sqrt(Ids)-Vgs plot
xfitmin = 1.90   # V, minimum Vgs for fitting
xfitmax = 10.0   # V, maximum Vgs for fitting


#===================================
# Utility functions
#===================================
def pfloat(value, defval=None):
    """
    Convert a string containing a floating point number (or a float) to float.

    Parameters:
        value : str or float
            Input that may contain a floating point number.
        defval : any, optional
            Value to return if conversion fails (default: None).

    Returns:
        float or defval if conversion fails.
    """
    if isinstance(value, float):
        return value
    m = re.search(r'([+\-eE\d\.]+)', value)
    if m:
        try:
            return float(m.group())
        except Exception:
            return defval
    return defval


def pint(value):
    """
    Convert a string to an integer.

    Parameters:
        value : str
            Input string.

    Returns:
        int or None if conversion fails.
    """
    try:
        return int(value)
    except Exception:
        return None


def getarg(position, defval=None):
    """
    Safely get command line argument by index.

    Parameters:
        position : int
            Index position of the argument.
        defval : any, optional
            Default value if the argument is not provided.

    Returns:
        str or defval if index is out of range.
    """
    try:
        return sys.argv[position]
    except Exception:
        return defval


def getfloatarg(position, defval=None):
    """
    Get command line argument as a float.

    Parameters:
        position : int
            Index position of the argument.
        defval : any, optional
            Default value if argument is not provided or conversion fails.

    Returns:
        float value of the argument.
    """
    return pfloat(getarg(position, defval))


def getintarg(position, defval=None):
    """
    Get command line argument as an integer.

    Parameters:
        position : int
            Index position of the argument.
        defval : any, optional
            Default value if argument is not provided or conversion fails.

    Returns:
        int value of the argument.
    """
    return pint(getarg(position, defval))


def usage():
    """
    Display usage information.
    """
    print("\nUsage:")
    print("  python {} (Vds0 xfitmin xfitmax)".format(sys.argv[0]))


def terminate():
    """
    Display usage and exit the program.
    """
    print("")
    usage()
    exit()


def savecsv(outfile, header, datalist):
    """
    Save data to a CSV file.

    Parameters:
        outfile : str
            Output CSV file name.
        header : list of str
            List of header strings.
        datalist : list of lists
            2D list where each sub-list represents a column.
    """
    try:
        print("Write to [{}]".format(outfile))
        f = open(outfile, 'w', newline='')
    except Exception:
        print("Error: Cannot write to [{}]".format(outfile))
        return
    else:
        fout = csv.writer(f)
        fout.writerow(header)
        # 行ごとにデータを書き出し
        for i in range(len(datalist[0])):
            row = [col[i] for col in datalist]
            fout.writerow(row)
        f.close()


def read_csv(fname):
    """
    Read a CSV file and extract header, x-data and y-data.

    Parameters:
        fname : str
            File name of the CSV.

    Returns:
        xlabel : str
            The label for x-axis.
        ylabels : list of str
            List of labels for y-data.
        x : list of float
            x-axis data.
        ylist : list of lists of float
            List containing each y-data as a list.
    """
    print("\nReading CSV file [{}]".format(fname))
    with open(fname) as f:
        fin = csv.reader(f)
        labels = next(fin)
        xlabel = labels[0]

        # 空文字で終わる場合はそれ以降はデータとして読み込まない
        ylabels = [label for label in labels[1:] if label != '']
        ny = len(ylabels)

        x = []
        ylist = [[] for _ in range(ny)]
        for row in fin:
            x.append(pfloat(row[0]))
            for i in range(ny):
                ylist[i].append(pfloat(row[i+1]))
    return xlabel, ylabels, x, ylist


def linear_regression(x, y, verbose=False):
    """
    Perform linear least-squares regression of y = a + b*x with error estimation.

    The function uses the standard formulas with degrees of freedom n-2.

    Parameters:
        x : list or array-like of float
            Independent variable data.
        y : list or array-like of float
            Dependent variable data.
        verbose : bool, optional
            If True, print detailed calculation results (default: False).

    Returns:
        a : float
            Intercept of the regression line.
        b : float
            Slope of the regression line.
        res : dict
            Dictionary containing:
                'SSE': Sum of squared errors.
                'sa': Standard error of the intercept.
                'sb': Standard error of the slope.
                's':  Standard deviation of the residuals.
                'r':  Correlation coefficient.
    """
    n = len(x)
    if n < 3:
        raise ValueError("Need at least 3 data points for linear regression (n-2 degrees of freedom).")

    avg_x = np.mean(x)
    avg_y = np.mean(y)

    # Calculate Sxx and Sxy using deviations from the mean
    Sxx = sum([(xi - avg_x) ** 2 for xi in x])
    Sxy = sum([(x[i] - avg_x) * (y[i] - avg_y) for i in range(n)])
    # Compute regression coefficients
    b = Sxy / Sxx
    a = avg_y - b * avg_x

    # Sum of squared errors and standard deviation (using degrees of freedom n-2)
    SSE = sum([(y[i] - a - b * x[i]) ** 2 for i in range(n)])
    s = sqrt(SSE / (n - 2))

    # Standard errors of intercept and slope
    sa = s * sqrt(1/n + (avg_x ** 2) / Sxx)
    sb = s / sqrt(Sxx)

    # Correlation coefficient r
    Syy = sum([(yi - avg_y) ** 2 for yi in y])
    r = Sxy / sqrt(Sxx * Syy)

    if verbose:
        print("\nLinear Regression Results:")
        print("n            =", n)
        print("avg_x        =", avg_x)
        print("avg_y        =", avg_y)
        print("Sxx          =", Sxx)
        print("Sxy          =", Sxy)
        print("SSE          =", SSE)
        print("Residual s   =", s)
        print("Intercept (a) =", a, " (sa =", sa, ")")
        print("Slope (b)    =", b, " (sb =", sb, ")")
        print("Correlation  =", r)

    res = {'SSE': SSE, 'sa': sa, 'sb': sb, 's': s, 'r': r}
    return a, b, res


def find_index(data, value, default=None):
    """
    Find the first index in 'data' where the element is greater than or equal to the specified value.
    
    Parameters:
        data : list of float
            Sorted list of numeric data.
        value : float
            The value to compare.
        default : any, optional
            Value to return if no suitable index is found.
    
    Returns:
        int or default value.
    """
    for i, d in enumerate(data):
        if d >= value - 1.0e-5:
            return i
    return default


#===================================
# Main processing function
#===================================
def main():
    """
    Main function to read CSV data, perform regression, and plot results.
    
    The function reads the CSV file, extracts Vgs and Ids data, computes sqrt(Ids),
    performs linear regression on a given fitting range, estimates threshold voltage and mobility,
    and plots multiple figures including transfer/output characteristics and error trends.
    """
    # Update parameters via command-line arguments if provided
    global Vds0, xfitmin, xfitmax
    if len(sys.argv) >= 2:
        Vds0 = float(sys.argv[1])
    if len(sys.argv) >= 3:
        xfitmin = float(sys.argv[2])
    if len(sys.argv) >= 4:
        xfitmax = float(sys.argv[3])
    Vds0    = getfloatarg(1, Vds0)
    xfitmin = getfloatarg(2, xfitmin)
    xfitmax = getfloatarg(3, xfitmax)

    # Calculate oxide capacitance per unit area (Cox)
    Cox = erg * e0 / dg  # [F/m^2]
    print("\nCox = {:12.6g} [F/m^2]".format(Cox))

    # Read CSV data
    xlabel, ylabels, Vgs, IdsVgs = read_csv(infile)
    nVgs = len(IdsVgs[0])
    nVds = len(ylabels)
    Vds = [pfloat(label) for label in ylabels]

    print("\nnVds =", nVds)
    print("nVgs =", nVgs)
    print("xlabel:", xlabel)
    print("ylabels:", ylabels)
    print("Vds:", Vds)

    # 転置して出力特性用の2次元配列を作成（各Vgsに対して複数のVdsデータ）
    IdsVds = np.empty([nVgs, nVds])
    for ig in range(nVgs):
        for id in range(nVds):
            IdsVds[ig][id] = IdsVgs[id][ig]

    # Vds0 に近いデータ番号を検索（sqrt(Ids)-Vgsプロット用）
    for i in range(nVds):
        if Vds[i] >= Vds0 - 1.0e-3:
            iVds = i
            break
    print("\nVds used: {} V (iVds = {})".format(Vds[iVds], iVds))

    # sqrt(Ids) データ作成
    sqrtIds = [sqrt(IdsVgs[iVds][ig]) for ig in range(nVgs)]

    # 全Vgsデータとsqrt(Ids)データをフィッティング用にコピー
    xfitall = Vgs[:]
    yfitall = sqrtIds[:]

    # フィッティング範囲の依存性確認
    print("\n=== Check the effect of fitting range ===")
    xecheck = [x for x in np.arange(xfitmin + 2.0, max(Vgs), 2.0)]
    ymu, yVth = [], []
    ySmu, ySVth = [], []
    yr, ysigma = [], []
    for current_xfitmax in xecheck:
        ixfitmin = find_index(xfitall, xfitmin)
        ixfitmax = find_index(xfitall, current_xfitmax, default=len(xfitall) - 1)
        print("\nFitting range: {} - {} V (indices {} - {})".format(xfitmin, current_xfitmax, ixfitmin, ixfitmax))

        # 線形回帰（verbose=False）
        a_fit, b_fit, regression_res = linear_regression(xfitall[ixfitmin:ixfitmax+1],
                                                         yfitall[ixfitmin:ixfitmax+1],
                                                         verbose=False)
        SSE = regression_res['SSE']
        s = regression_res['s']
        sa = regression_res['sa']
        sb = regression_res['sb']
        r = regression_res['r']
        # 閾値電圧 Vth と移動度 mu の計算（誤差伝搬は簡易的な計算）
        Vth = -a_fit / b_fit
        grad = b_fit
        mu = grad**2 / (W * Cox / 2.0 / L)
        # 誤差伝搬（詳細な計算は用途に応じて見直してください）
        SVth = sqrt((sa / b_fit)**2 + (sb * a_fit / (b_fit**2))**2)
        Smu = 2.0 * mu * sb / b_fit

        yVth.append(Vth)
        ymu.append(mu)
        ySVth.append(SVth)
        ySmu.append(Smu)
        yr.append(r)
        ysigma.append(s)

        print("Residual SSE = {:8.6g}".format(SSE))
        print("Standard deviation (s) = {:8.6g}".format(s))
        print("Correlation coefficient r = {:8.6g}".format(r))
        print("Vth = {:6.4g} (±{:12.4g}) V".format(Vth, SVth))
        print("mu_sat = {:12.4g} (±{:8.4g}) cm^2/Vs".format(1.0e4 * mu, 1.0e4 * Smu))

    # 固定フィッティング範囲での最小二乗フィッティング
    print("\n=== Fitting by the given fitting range ===")
    ixfitmin = find_index(xfitall, xfitmin)
    ixfitmax = find_index(xfitall, xfitmax, default=len(xfitall) - 1)
    xfit = xfitall[ixfitmin:ixfitmax+1]
    yfit = yfitall[ixfitmin:ixfitmax+1]

    print("\nLeast squares fitting:")
    print("Vgs range: {} - {} V (indices {} - {})".format(xfitmin, xfitmax, ixfitmin, ixfitmax))
    print("Vgs =", xfit)
    print("sqrt(Ids) =", yfit)

    a_fit, b_fit, regression_res = linear_regression(xfit, yfit, verbose=False)
    sa = regression_res['sa']
    sb = regression_res['sb']
    r  = regression_res['r']
    Vth = -a_fit / b_fit
    grad = b_fit
    mu = grad**2 / (W * Cox / 2.0 / L)
    SVth = sqrt((sa / b_fit)**2 + (sb * a_fit / (b_fit**2))**2)
    Smu = 2.0 * mu * sb / b_fit

    print("\ny = a + bx")
    print("  a = {:12.6g} (±{:12.4g})".format(a_fit, sa))
    print("  b = {:12.6g} (±{:12.4g})".format(b_fit, sb))
    print("  r = {:8.6g}".format(r))
    print("Vth = {:6.4g} (±{:12.4g}) V".format(Vth, SVth))
    print("d(sqrt(Ids))/dVgs = {:12.4g} (±{:12.4g})".format(grad, sb))
    print("mu_sat = {:12.4g} m^2/Vs  (±{:8.4g} m^2/Vs) = {:12.4g} cm^2/Vs  (±{:8.4g} cm^2/Vs)"
          .format(mu, Smu, 1.0e4 * mu, 1.0e4 * Smu))

    # プロット用の直線計算
    xcal = [xfitmin - 0.5, max(Vgs)]
    ycal = [a_fit + b_fit * xx for xx in xcal]

    # プロット設定
    fig = plt.figure(figsize=(12, 8))
    ax1 = fig.add_subplot(3, 3, 1)
    ax2 = fig.add_subplot(3, 3, 2)
    ax3 = fig.add_subplot(3, 3, 3)
    ax4 = fig.add_subplot(3, 3, 4)
    ax5 = fig.add_subplot(3, 3, 5)
    ax6 = fig.add_subplot(3, 3, 6)
    ax7 = fig.add_subplot(3, 3, 7)

    # Transfer characteristics (Ids-Vgs)
    for id in range(nVds):
        ax1.plot(Vgs, IdsVgs[id], linewidth=0.5, marker='o', markersize=0.5,
                 label='Vds={} V'.format(Vds[id]))
    ax1.set_xlabel('Vgs (V)')
    ax1.set_ylabel("Ids (A)")
    ax1.set_yscale('log')

    # Output characteristics (Ids-Vds)
    nskip = int(nVgs / 20.0 + 1.0e-6)
    for ig in range(0, nVgs, nskip):
        if max(IdsVds[ig]) < 1.0e-7:
            continue
        ax2.plot(Vds, IdsVds[ig], linewidth=0.5, marker='o', markersize=0.5,
                 label='Vgs={} V'.format(Vgs[ig]))
    ax2.set_xlabel('Vds (V)')
    ax2.set_ylabel("Ids (A)")
    ax2.legend(loc='upper left', fontsize=6)

    # sqrt(Ids)-Vgs plot with fit line
    ax3.plot(Vgs, sqrtIds, linestyle='none', marker='o', markersize=0.5)
    ax3.plot(xcal, ycal, linestyle='-', linewidth=0.5)
    ax3.set_xlabel('Vgs (V)')
    ax3.set_ylabel("sqrt(Ids) (A^(1/2))")

    # Mobility vs. fitting range plot with error bars
    ax4.errorbar(xecheck, ymu, yerr=ySmu, capsize=3.0, fmt='o', markersize=3.0,
                 ecolor='b', markeredgecolor='b', color='w')
    ax4.set_xlabel("Vgs max for fitting (V)")
    ax4.set_ylabel('Mobility (m^2/Vs)')

    # Threshold voltage vs. fitting range plot with error bars
    ax5.errorbar(xecheck, yVth, yerr=ySVth, capsize=3.0, fmt='o', markersize=3.0,
                 ecolor='b', markeredgecolor='b', color='w')
    ax5.set_xlabel("Vgs max for fitting (V)")
    ax5.set_ylabel('Vth (V)')

    # Correlation coefficient vs. fitting range
    ax6.plot(xecheck, yr, marker='o')
    ax6.set_xlabel("Vgs max for fitting (V)")
    ax6.set_ylabel('Correlation Coefficient')

    # Residual standard deviation vs. fitting range
    ax7.plot(xecheck, ysigma, marker='o')
    ax7.set_xlabel("Vgs max for fitting (V)")
    ax7.set_ylabel('Residual Std Dev (s)')

    plt.tight_layout()
    plt.pause(0.1)
    input("\nPress ENTER to exit>>")
    terminate()


if __name__ == '__main__':
    main()