#!/usr/bin/env python3
"""Analyze TFT transfer characteristics from Id-Vg Excel files.

This script reads two Excel files, one for the linear-region transfer curve
and one for the saturation-region transfer curve, and calculates typical TFT
parameters such as Vth, mobility, SS, and Ion.

The default input/output filenames are kept compatible with the original
script:

    IdVg-Vd0.1.xlsx
    IdVg-Vd10.xlsx

Default output files:

    IdVg-Vd10_analyze.csv
    IdVg-Vd0.1_analyze.csv
    output.csv
    output_log.csv
    rootIdVg.png
    IdVg_LIN.png
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from dataclasses import dataclass
from typing import Optional

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.linear_model import LinearRegression


# -----------------------------------------------------------------------------
# Default analysis parameters kept from the original script
# -----------------------------------------------------------------------------
DEFAULT_COX = 2.24e-8
DEFAULT_L = 50.0e-6
DEFAULT_W = 300.0e-6
DEFAULT_TARGET_CURRENT = 1.0e-10
DEFAULT_IOFF_THRESHOLD = 1.0e-14
DEFAULT_ROLLING_WINDOW = 7


@dataclass(frozen=True)
class AnalysisConfig:
    """Configuration values used for TFT parameter extraction.

    Attributes:
        cox: Gate capacitance per unit area. The unit must be consistent with L,
            W, Id, and Vd. In the original script this was 2.24e-8.
        channel_length: Channel length L.
        channel_width: Channel width W.
        target_current: Current threshold used to evaluate SS and Von.
        ioff_threshold: Upper current threshold used for Ioff averaging.
        rolling_window: Window size for rolling average of the local slope.
    """

    cox: float = DEFAULT_COX
    channel_length: float = DEFAULT_L
    channel_width: float = DEFAULT_W
    target_current: float = DEFAULT_TARGET_CURRENT
    ioff_threshold: float = DEFAULT_IOFF_THRESHOLD
    rolling_window: int = DEFAULT_ROLLING_WINDOW


@dataclass
class IdVgResult:
    """Summary values calculated from an Id-Vg curve."""

    dataframe: pd.DataFrame
    vth: float
    max_slope: float
    vg_at_max_slope: float
    mobility: float
    ss: float
    ion: float
    ioff: Optional[float]
    ion_ioff_ratio: Optional[float]
    von: float


@dataclass
class IdVdResult:
    """Summary values calculated from an Id-Vd curve."""

    dataframe: pd.DataFrame
    mueff: float


# -----------------------------------------------------------------------------
# Input and validation
# -----------------------------------------------------------------------------
def read_excel_file(path: str, label: str) -> pd.DataFrame:
    """Read an Excel file and add a clear error message on failure.

    Args:
        path: Path to the Excel file.
        label: Human-readable name used in error messages.

    Returns:
        DataFrame loaded from the Excel file.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If pandas cannot read the file.
    """

    if not os.path.exists(path):
        raise FileNotFoundError(f"{label} file not found: {path}")

    try:
        return pd.read_excel(path)
    except Exception as exc:  # pandas may raise different exceptions by engine/version
        raise ValueError(f"Failed to read {label} Excel file: {path}") from exc


def validate_required_columns(df: pd.DataFrame, required_columns: list[str], label: str) -> None:
    """Validate that a DataFrame contains required columns.

    Args:
        df: DataFrame to validate.
        required_columns: Column names required by the analysis.
        label: Human-readable name used in error messages.

    Raises:
        ValueError: If one or more required columns are missing.
    """

    missing = [col for col in required_columns if col not in df.columns]
    if missing:
        raise ValueError(
            f"{label} data is missing required columns: {missing}. "
            f"Available columns: {list(df.columns)}"
        )


# -----------------------------------------------------------------------------
# Numerical helper functions
# -----------------------------------------------------------------------------
def calculate_local_slope_by_linear_fit(
    df: pd.DataFrame,
    x_series: pd.Series,
    y_series: pd.Series,
    half_window: int = 1,
) -> list[float]:
    """Calculate local slopes by fitting y = a*x + b around each data point.

    This function replaces the original ``slope3data`` and ``slope7data`` style
    functions. With ``half_window=1``, it uses approximately three points,
    matching the original ``slope3data`` behavior as closely as possible.

    Args:
        df: Source DataFrame. Only its length is used.
        x_series: X values.
        y_series: Y values.
        half_window: Number of neighboring points on each side.

    Returns:
        List of local slopes. Points that cannot be fitted are returned as NaN.
    """

    slopes: list[float] = []

    for index in range(len(df)):
        x_window = x_series[index - half_window : index + half_window + 1]
        y_window = y_series[index - half_window : index + half_window + 1]

        valid = pd.concat([x_window, y_window], axis=1).dropna()
        if len(valid) < 2:
            slopes.append(np.nan)
            continue

        try:
            x = valid.iloc[:, 0].to_numpy().reshape(-1, 1)
            y = valid.iloc[:, 1].to_numpy()
            model = LinearRegression().fit(x, y)
            slopes.append(float(model.coef_[0]))
        except ValueError as exc:
            logging.debug("Local slope fitting failed at index %s: %s", index, exc)
            slopes.append(np.nan)

    return slopes


def intercept_x(x: float, y: float, slope: float) -> float:
    """Calculate the x-intercept of the line y_line = slope * (x_line - x0).

    Args:
        x: X coordinate of a point on the line.
        y: Y coordinate of a point on the line.
        slope: Slope of the line.

    Returns:
        X-intercept value.
    """

    if pd.isna(slope) or slope == 0:
        return np.nan
    return float(x - y / slope)


def calculate_vth_by_max_slope(
    vg_series: pd.Series,
    y_series: pd.Series,
    slope_series: pd.Series,
) -> tuple[float, float, float]:
    """Calculate threshold voltage using the maximum-slope tangent method.

    Args:
        vg_series: Gate voltage values.
        y_series: Id or sqrt(Id), depending on the analysis mode.
        slope_series: Local slope values.

    Returns:
        Tuple of ``(Vth, max_slope, Vg_at_max_slope)``.

    Raises:
        ValueError: If no valid slope can be found.
    """

    valid_slope = slope_series.replace([np.inf, -np.inf], np.nan).dropna()
    if valid_slope.empty:
        raise ValueError("No valid slope values are available for Vth calculation.")

    idx = valid_slope.idxmax()
    max_slope = float(valid_slope.loc[idx])
    x = float(vg_series.loc[idx])
    y = float(y_series.loc[idx])
    vth = intercept_x(x, y, max_slope)
    return vth, max_slope, x


def calculate_ion_ioff(id_series: pd.Series, ioff_threshold: float) -> tuple[float, Optional[float], Optional[float]]:
    """Calculate Ion, Ioff, and Ion/Ioff ratio.

    Args:
        id_series: Drain current values. Positive values are expected.
        ioff_threshold: Data points satisfying ``0 < Id < ioff_threshold`` are
            averaged to estimate Ioff.

    Returns:
        Tuple of ``(Ion, Ioff, Ion/Ioff)``. If Ioff cannot be estimated,
        ``Ioff`` and ``Ion/Ioff`` are returned as None.
    """

    valid_id = pd.to_numeric(id_series, errors="coerce").replace([np.inf, -np.inf], np.nan)
    ion = float(valid_id.max()) if not valid_id.dropna().empty else np.nan

    ioff_data = valid_id[(valid_id > 0) & (valid_id < ioff_threshold)].dropna()
    if ioff_data.empty:
        return ion, None, None

    ioff = float(ioff_data.mean())
    ratio = float(ion / ioff) if ioff != 0 else None
    return ion, ioff, ratio


def first_index_at_or_above(series: pd.Series, threshold: float) -> Optional[int]:
    """Return the first index where the series value is at least threshold.

    Args:
        series: Numeric series.
        threshold: Threshold value.

    Returns:
        First matching index, or None if no point satisfies the condition.
    """

    matched = series.index[series >= threshold]
    if len(matched) == 0:
        return None
    return matched.min()


def safe_inverse(value: float) -> float:
    """Return 1/value, or NaN if the inverse is not well-defined."""

    if pd.isna(value) or value == 0:
        return np.nan
    return float(1.0 / value)


# -----------------------------------------------------------------------------
# TFT analysis functions
# -----------------------------------------------------------------------------
def analyze_idvg_sat(df: pd.DataFrame, config: AnalysisConfig) -> IdVgResult:
    """Analyze saturation-region Id-Vg data.

    The calculation follows the original script:

    - Convert negative/non-positive Id to NaN for log and sqrt operations.
    - Calculate log10(Id) and sqrt(Id).
    - Calculate local slopes by three-point linear fitting.
    - Smooth sqrt(Id) slope by centered rolling average.
    - Calculate saturation mobility from the smoothed sqrt(Id) slope.
    - Calculate Vth by maximum-slope tangent method.
    - Calculate SS at the first point where Id >= target_current.

    Args:
        df: Input DataFrame. Required columns: ``Vg`` and ``Id``.
        config: Analysis configuration.

    Returns:
        IdVgResult containing the processed DataFrame and summary values.
    """

    validate_required_columns(df, ["Vg", "Id"], "SAT Id-Vg")
    df = df.copy()

    df["Id"] = pd.to_numeric(df["Id"], errors="coerce")
    df["Id_rm"] = np.where(df["Id"] > 0, df["Id"], np.nan)
    df["logId"] = np.log10(df["Id_rm"])
    df["rootId"] = df["Id_rm"] ** 0.5

    df["slope logId"] = calculate_local_slope_by_linear_fit(df, df["Vg"], df["logId"], half_window=1)
    df["slope rootId"] = calculate_local_slope_by_linear_fit(df, df["Vg"], df["rootId"], half_window=1)
    df["slope logId"] = df["slope logId"].replace([np.inf, -np.inf], np.nan)
    df["slope rootId"] = df["slope rootId"].replace([np.inf, -np.inf], np.nan)
    df["slope rootId smooth"] = df["slope rootId"].rolling(config.rolling_window, center=True).mean()

    df["muSAT"] = (
        2.0
        * config.channel_length
        / config.channel_width
        / config.cox
        * (df["slope rootId smooth"] ** 2)
    )
    df["muSAT"] = df["muSAT"].replace([np.inf, -np.inf], np.nan)

    vth, max_slope, vg_at_max_slope = calculate_vth_by_max_slope(
        df["Vg"], df["rootId"], df["slope rootId smooth"]
    )
    mobility = float(df["muSAT"].max()) if not df["muSAT"].dropna().empty else np.nan

    idx = first_index_at_or_above(df["Id_rm"], config.target_current)
    if idx is None:
        logging.warning("SAT: Id never reaches target current %.3e A. SS and Von are set to NaN.", config.target_current)
        ss = np.nan
        von = np.nan
    else:
        ss = safe_inverse(df.loc[idx, "slope logId"])
        von = float(df.loc[idx, "Vg"])

    ion, ioff, ratio = calculate_ion_ioff(df["Id_rm"], config.ioff_threshold)

    return IdVgResult(
        dataframe=df,
        vth=vth,
        max_slope=max_slope,
        vg_at_max_slope=vg_at_max_slope,
        mobility=mobility,
        ss=ss,
        ion=ion,
        ioff=ioff,
        ion_ioff_ratio=ratio,
        von=von,
    )


def analyze_idvg_lin(df: pd.DataFrame, config: AnalysisConfig) -> IdVgResult:
    """Analyze linear-region Id-Vg data.

    The calculation follows the original script:

    - Convert negative/non-positive Id to NaN.
    - Calculate local slopes of Id and log10(Id).
    - Smooth Id slope by centered rolling average.
    - Calculate field-effect mobility from the smoothed Id slope.
    - Calculate Vth by maximum-slope tangent method.
    - Calculate SS at the first point where Id >= target_current.

    Args:
        df: Input DataFrame. Required columns: ``Vg``, ``Vd``, and ``Id``.
        config: Analysis configuration.

    Returns:
        IdVgResult containing the processed DataFrame and summary values.
    """

    validate_required_columns(df, ["Vg", "Vd", "Id"], "LIN Id-Vg")
    df = df.copy()

    df["Id"] = pd.to_numeric(df["Id"], errors="coerce")
    df["Id_rm"] = np.where(df["Id"] > 0, df["Id"], np.nan)
    df["logId"] = np.log10(df["Id_rm"])

    df["slope Id"] = calculate_local_slope_by_linear_fit(df, df["Vg"], df["Id_rm"], half_window=1)
    df["slope logId"] = calculate_local_slope_by_linear_fit(df, df["Vg"], df["logId"], half_window=1)
    df["slope Id"] = df["slope Id"].replace([np.inf, -np.inf], np.nan)
    df["slope Id smooth"] = df["slope Id"].rolling(config.rolling_window, center=True).mean()
    df["slope logId"] = df["slope logId"].replace([np.inf, -np.inf], np.nan)

    df["muFE"] = (
        config.channel_length
        / config.channel_width
        / config.cox
        / df["Vd"]
        * df["slope Id smooth"]
    )
    df["muFE"] = df["muFE"].replace([np.inf, -np.inf], np.nan)

    vth, max_slope, vg_at_max_slope = calculate_vth_by_max_slope(
        df["Vg"], df["Id_rm"], df["slope Id smooth"]
    )
    mobility = float(df["muFE"].max()) if not df["muFE"].dropna().empty else np.nan

    idx = first_index_at_or_above(df["Id_rm"], config.target_current)
    if idx is None:
        logging.warning("LIN: Id never reaches target current %.3e A. SS and Von are set to NaN.", config.target_current)
        ss = np.nan
        von = np.nan
    else:
        ss = safe_inverse(df.loc[idx, "slope logId"])
        von = float(df.loc[idx, "Vg"])

    ion, ioff, ratio = calculate_ion_ioff(df["Id_rm"], config.ioff_threshold)

    return IdVgResult(
        dataframe=df,
        vth=vth,
        max_slope=max_slope,
        vg_at_max_slope=vg_at_max_slope,
        mobility=mobility,
        ss=ss,
        ion=ion,
        ioff=ioff,
        ion_ioff_ratio=ratio,
        von=von,
    )


def analyze_idvd_lin(df: pd.DataFrame, vth: float, config: AnalysisConfig) -> IdVdResult:
    """Analyze linear-region Id-Vd data.

    This function is kept because the original script contained ``analyze_IdVd_LIN``,
    although it was not called by the main process.

    Args:
        df: Input DataFrame. Required columns: ``Vd``, ``Vg``, and ``Id``.
        vth: Threshold voltage used for mueff calculation.
        config: Analysis configuration.

    Returns:
        IdVdResult containing the processed DataFrame and maximum mueff.
    """

    validate_required_columns(df, ["Vd", "Vg", "Id"], "LIN Id-Vd")
    df = df.copy()

    df["Id"] = pd.to_numeric(df["Id"], errors="coerce")
    df["Id_rm"] = np.where(df["Id"] > 0, df["Id"], np.nan)
    df = df.sort_values(by="Vd")
    df["slope Id"] = calculate_local_slope_by_linear_fit(df, df["Vd"], df["Id_rm"], half_window=1)
    df["mueff"] = (
        config.channel_length
        / config.channel_width
        / config.cox
        / (df["Vg"] - vth)
        * df["slope Id"]
    )
    df["mueff"] = df["mueff"].replace([np.inf, -np.inf], np.nan)
    mueff = float(df["mueff"].max()) if not df["mueff"].dropna().empty else np.nan

    return IdVdResult(dataframe=df, mueff=mueff)


# -----------------------------------------------------------------------------
# Plot functions
# -----------------------------------------------------------------------------
def plot_idvg_lin(df: pd.DataFrame, vth: float, max_slope: float, output_path: str) -> None:
    """Save the linear-region Id-Vg plot with maximum-slope tangent line."""

    plt.figure()
    plt.scatter(df["Vg"], df["Id"], label="Vd=0.1V", s=4)
    df_line = max_slope * (df["Vg"] - vth)
    plt.plot(df["Vg"], df_line, color="red")
    plt.xlabel("Vg(V)")
    plt.ylabel("Id(A)")
    plt.ylim(0, 1.1 * df["Id"].max())
    plt.legend()
    plt.savefig(output_path)
    plt.close()


def plot_root_idvg(df: pd.DataFrame, vth: float, max_slope: float, output_path: str) -> None:
    """Save the saturation-region sqrt(Id)-Vg plot with tangent line."""

    plt.figure()
    plt.scatter(df["Vg"], df["rootId"], label="Vd=10V", s=4)
    df_line = max_slope * (df["Vg"] - vth)
    plt.plot(df["Vg"], df_line, color="red")
    plt.xlabel("Vg(V)")
    plt.ylabel("Id^0.5(A^0.5)")
    plt.ylim(0, 1.1 * df["rootId"].max())
    plt.legend()
    plt.savefig(output_path)
    plt.close()


# -----------------------------------------------------------------------------
# Output functions
# -----------------------------------------------------------------------------
def output_path(output_dir: str, filename: str) -> str:
    """Return a path inside the output directory."""

    return os.path.join(output_dir, filename)


def write_summary_outputs(sat: IdVgResult, lin: IdVgResult, output_dir: str) -> None:
    """Write summary CSV files compatible with the original script.

    This function intentionally keeps the original output column names for
    compatibility, including ``output_log.csv`` where SS and Ion values are
    converted to log10 without changing the column names.
    """

    df_summary = pd.DataFrame()

    df_summary["Vth_SAT"] = [sat.vth]
    # df_summary["Von_SAT"] = [sat.von]
    df_summary["muSAT"] = [sat.mobility]
    # df_summary["Vgmaxslope_SAT"] = [sat.vg_at_max_slope]
    df_summary["SS_SAT"] = [sat.ss]
    df_summary["Ion_SAT"] = [sat.ion]

    df_summary["Vth_LIN"] = [lin.vth]
    # df_summary["Von_LIN"] = [lin.von]
    df_summary["muFE"] = [lin.mobility]
    # df_summary["Vgmaxslope_LIN"] = [lin.vg_at_max_slope]
    df_summary["SS_LIN"] = [lin.ss]
    df_summary["Ion_LIN"] = [lin.ion]

    df_summary.to_csv(output_path(output_dir, "output.csv"), index=False)

    df_log = df_summary.copy()
    for column in ["SS_SAT", "Ion_SAT", "SS_LIN", "Ion_LIN"]:
        df_log[column] = np.log10(df_log[column])
    df_log.to_csv(output_path(output_dir, "output_log.csv"), index=False)


# -----------------------------------------------------------------------------
# Command line interface
# -----------------------------------------------------------------------------
def parse_args(argv: Optional[list[str]] = None) -> argparse.Namespace:
    """Parse command-line arguments."""

    parser = argparse.ArgumentParser(
        description="Analyze TFT Id-Vg data and calculate Vth, mobility, SS, and Ion."
    )
    parser.add_argument(
        "--lin",
        default="IdVg-Vd0.1.xlsx",
        help="Excel file for linear-region Id-Vg data. Default: %(default)s",
    )
    parser.add_argument(
        "--sat",
        default="IdVg-Vd10.xlsx",
        help="Excel file for saturation-region Id-Vg data. Default: %(default)s",
    )
    parser.add_argument(
        "--outdir",
        default=".",
        help="Output directory. Default: current directory",
    )
    parser.add_argument(
        "--cox",
        type=float,
        default=DEFAULT_COX,
        help="Gate capacitance per unit area. Default: %(default).3e",
    )
    parser.add_argument(
        "--length",
        type=float,
        default=DEFAULT_L,
        help="Channel length L. Default: %(default).3e",
    )
    parser.add_argument(
        "--width",
        type=float,
        default=DEFAULT_W,
        help="Channel width W. Default: %(default).3e",
    )
    parser.add_argument(
        "--target-current",
        type=float,
        default=DEFAULT_TARGET_CURRENT,
        help="Current threshold used for SS and Von. Default: %(default).3e",
    )
    parser.add_argument(
        "--ioff-threshold",
        type=float,
        default=DEFAULT_IOFF_THRESHOLD,
        help="Upper current threshold used for Ioff averaging. Default: %(default).3e",
    )
    parser.add_argument(
        "--rolling-window",
        type=int,
        default=DEFAULT_ROLLING_WINDOW,
        help="Centered rolling-average window for slope smoothing. Default: %(default)s",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Show detailed traceback on error.",
    )
    return parser.parse_args(argv)


def validate_config(config: AnalysisConfig) -> None:
    """Validate command-line configuration values."""

    if config.cox == 0:
        raise ValueError("cox must not be zero.")
    if config.channel_width == 0:
        raise ValueError("channel width must not be zero.")
    if config.rolling_window < 1:
        raise ValueError("rolling-window must be at least 1.")
    if config.target_current <= 0:
        raise ValueError("target-current must be positive.")
    if config.ioff_threshold <= 0:
        raise ValueError("ioff-threshold must be positive.")


def run(args: argparse.Namespace) -> None:
    """Run the complete analysis workflow."""

    config = AnalysisConfig(
        cox=args.cox,
        channel_length=args.length,
        channel_width=args.width,
        target_current=args.target_current,
        ioff_threshold=args.ioff_threshold,
        rolling_window=args.rolling_window,
    )
    validate_config(config)

    os.makedirs(args.outdir, exist_ok=True)

    df_sat_input = read_excel_file(args.sat, "SAT Id-Vg")
    df_lin_input = read_excel_file(args.lin, "LIN Id-Vg")

    sat = analyze_idvg_sat(df_sat_input, config)
    sat.dataframe.to_csv(output_path(args.outdir, "IdVg-Vd10_analyze.csv"), index=False)
    plot_root_idvg(
        sat.dataframe,
        sat.vth,
        sat.max_slope,
        output_path(args.outdir, "rootIdVg.png"),
    )

    lin = analyze_idvg_lin(df_lin_input, config)
    lin.dataframe.to_csv(output_path(args.outdir, "IdVg-Vd0.1_analyze.csv"), index=False)
    plot_idvg_lin(
        lin.dataframe,
        lin.vth,
        lin.max_slope,
        output_path(args.outdir, "IdVg_LIN.png"),
    )

    write_summary_outputs(sat, lin, args.outdir)

    print("Analysis finished.")
    print(f"Output directory: {args.outdir}")
    print(f"Vth_SAT = {sat.vth}")
    print(f"muSAT   = {sat.mobility}")
    print(f"SS_SAT  = {sat.ss}")
    print(f"Ion_SAT = {sat.ion}")
    print(f"Vth_LIN = {lin.vth}")
    print(f"muFE    = {lin.mobility}")
    print(f"SS_LIN  = {lin.ss}")
    print(f"Ion_LIN = {lin.ion}")


def main(argv: Optional[list[str]] = None) -> int:
    """Entry point for command-line execution."""

    args = parse_args(argv)
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s: %(message)s",
    )

    try:
        run(args)
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        if args.verbose:
            raise
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())