"""Python code generator for NLSQ GUI.

This module generates standalone Python scripts that reproduce curve fitting
results from the GUI session. The generated scripts include all necessary
imports, data definitions, model functions, and curve_fit calls.

Functions
---------
generate_fit_script
    Generate a complete Python script from session state and fit result.
"""

from typing import Any

import numpy as np


def generate_fit_script(
    state: Any,
    result: Any,
) -> str:
    """Generate a standalone Python script that reproduces the fit.

    The generated script includes:
    - Import statements for nlsq and numpy
    - Data array definitions
    - Model function definition
    - curve_fit call with the same parameters
    - Result printing

    Parameters
    ----------
    state : SessionState
        The GUI session state containing data, model, and configuration.
    result : CurveFitResult
        The fitting result with fitted parameters.

    Returns
    -------
    str
        Complete Python script as a string.

    Examples
    --------
    >>> code = generate_fit_script(state, result)
    >>> # Run the generated code to reproduce the fit
    """
    lines: list[str] = []

    # Header comment
    lines.extend(
        [
            '"""',
            "NLSQ Curve Fitting Script",
            "",
            "This script was generated by the NLSQ GUI to reproduce a curve fit.",
            "It can be run standalone to perform the same fitting operation.",
            '"""',
            "",
        ]
    )

    # Imports
    lines.extend(
        [
            "# =============================================================================",
            "# Imports",
            "# =============================================================================",
            "",
            "import numpy as np",
            "import jax.numpy as jnp",
            "from nlsq import curve_fit",
            "",
            "",
        ]
    )

    # Data definition
    lines.extend(_generate_data_section(state))

    # Model definition
    lines.extend(_generate_model_section(state))

    # Initial guess and bounds
    lines.extend(_generate_parameters_section(state, result))

    # Fitting call
    lines.extend(_generate_fit_section(state))

    # Results printing
    lines.extend(_generate_results_section())

    return "\n".join(lines)


def _generate_data_section(state: Any) -> list[str]:
    """Generate the data definition section."""
    lines: list[str] = [
        "# =============================================================================",
        "# Data",
        "# =============================================================================",
        "",
    ]

    xdata = getattr(state, "xdata", None)
    ydata = getattr(state, "ydata", None)
    sigma = getattr(state, "sigma", None)

    if xdata is not None:
        xdata_str = _format_array(np.asarray(xdata))
        lines.append(f"xdata = np.array({xdata_str})")
    else:
        lines.append("xdata = np.array([])  # No data loaded")

    lines.append("")

    if ydata is not None:
        ydata_str = _format_array(np.asarray(ydata))
        lines.append(f"ydata = np.array({ydata_str})")
    else:
        lines.append("ydata = np.array([])  # No data loaded")

    lines.append("")

    if sigma is not None:
        sigma_str = _format_array(np.asarray(sigma))
        lines.append(f"sigma = np.array({sigma_str})")
    else:
        lines.append("sigma = None  # No uncertainties provided")

    lines.extend(["", ""])

    return lines


def _generate_model_section(state: Any) -> list[str]:
    """Generate the model function definition section."""
    lines: list[str] = [
        "# =============================================================================",
        "# Model Function",
        "# =============================================================================",
        "",
    ]

    model_type = getattr(state, "model_type", "builtin")
    model_name = getattr(state, "model_name", "exponential_decay")

    if model_type == "builtin":
        lines.extend(_get_builtin_model_code(model_name))
    elif model_type == "polynomial":
        degree = getattr(state, "polynomial_degree", 3)
        lines.extend(_get_polynomial_model_code(degree))
    elif model_type == "custom":
        custom_code = getattr(state, "custom_code", "")
        if custom_code:
            lines.append("# Custom model function")
            lines.append(custom_code)
        else:
            lines.append("# Custom model function (not defined)")
            lines.append("def model(x, *params):")
            lines.append("    return x  # Placeholder")
    else:
        lines.append("# Unknown model type")
        lines.append("def model(x, *params):")
        lines.append("    return x  # Placeholder")

    lines.extend(["", ""])

    return lines


def _generate_parameters_section(state: Any, result: Any) -> list[str]:
    """Generate the initial guess and bounds section."""
    lines: list[str] = [
        "# =============================================================================",
        "# Initial Guess and Bounds",
        "# =============================================================================",
        "",
    ]

    # Use fitted parameters as p0 for reproducibility
    popt = getattr(result, "popt", None)
    if popt is not None:
        p0_str = _format_array(np.asarray(popt))
        lines.append("# Using fitted values as initial guess for reproducibility")
        lines.append(f"p0 = {p0_str}")
    else:
        p0 = getattr(state, "p0", None)
        if p0 is not None:
            lines.append(f"p0 = {list(p0)}")
        else:
            lines.append("p0 = None  # Auto-estimate initial guess")

    lines.append("")

    bounds = getattr(state, "bounds", None)
    if bounds is not None:
        lower = list(bounds[0]) if bounds[0] is not None else None
        upper = list(bounds[1]) if bounds[1] is not None else None
        if lower is not None and upper is not None:
            lines.append(f"bounds = ({lower}, {upper})")
        else:
            lines.append("bounds = (-np.inf, np.inf)")
    else:
        lines.append("bounds = (-np.inf, np.inf)")

    lines.extend(["", ""])

    return lines


def _generate_fit_section(state: Any) -> list[str]:
    """Generate the curve_fit call section."""
    lines: list[str] = [
        "# =============================================================================",
        "# Perform Curve Fit",
        "# =============================================================================",
        "",
    ]

    # Build fit call with configuration
    gtol = getattr(state, "gtol", 1e-8)
    ftol = getattr(state, "ftol", 1e-8)
    xtol = getattr(state, "xtol", 1e-8)

    lines.extend(
        [
            "# Run the fit",
            "result = curve_fit(",
            "    model,",
            "    xdata,",
            "    ydata,",
            "    p0=p0,",
            "    sigma=sigma,",
            "    bounds=bounds,",
            f"    gtol={gtol},",
            f"    ftol={ftol},",
            f"    xtol={xtol},",
            ")",
            "",
            "",
        ]
    )

    return lines


def _generate_results_section() -> list[str]:
    """Generate the results printing section."""
    lines: list[str] = [
        "# =============================================================================",
        "# Print Results",
        "# =============================================================================",
        "",
        'print("Fitting Results")',
        'print("=" * 50)',
        "print()",
        "",
        "# Fitted parameters",
        'print("Fitted Parameters:")',
        "for i, (val, err) in enumerate(zip(result.popt, result.perr)):",
        '    print(f"  p{i}: {val:.6g} +/- {err:.6g}")',
        "print()",
        "",
        "# Fit quality metrics",
        'print("Fit Statistics:")',
        'print(f"  R-squared: {result.r_squared:.6f}")',
        'print(f"  RMSE: {result.rmse:.6g}")',
        'print(f"  MAE: {result.mae:.6g}")',
        "print()",
        "",
        "# Convergence info",
        'print("Convergence:")',
        'print(f"  Success: {result.success}")',
        'print(f"  Message: {result.message}")',
        'print(f"  Function evaluations: {result.nfev}")',
    ]

    return lines


def _format_array(arr: np.ndarray, max_elements: int = 20) -> str:
    """Format a numpy array as a Python list literal."""
    flat = arr.flatten()

    if len(flat) <= max_elements:
        # Format all elements
        elements = [f"{x:.6g}" for x in flat]
        return "[" + ", ".join(elements) + "]"
    else:
        # Too many elements, show first and last
        first = [f"{x:.6g}" for x in flat[:5]]
        last = [f"{x:.6g}" for x in flat[-5:]]
        return "[" + ", ".join(first) + ", ..., " + ", ".join(last) + "]"


def _get_builtin_model_code(model_name: str) -> list[str]:
    """Get Python code for a built-in model."""
    models = {
        "exponential_decay": [
            "def model(x, a, b, c):",
            '    """Exponential decay: a * exp(-b * x) + c"""',
            "    return a * jnp.exp(-b * x) + c",
        ],
        "exponential_growth": [
            "def model(x, a, b, c):",
            '    """Exponential growth: a * exp(b * x) + c"""',
            "    return a * jnp.exp(b * x) + c",
        ],
        "gaussian": [
            "def model(x, a, mu, sigma):",
            '    """Gaussian: a * exp(-((x - mu) / sigma)^2 / 2)"""',
            "    return a * jnp.exp(-((x - mu) / sigma) ** 2 / 2)",
        ],
        "lorentzian": [
            "def model(x, a, x0, gamma):",
            '    """Lorentzian: a / (1 + ((x - x0) / gamma)^2)"""',
            "    return a / (1 + ((x - x0) / gamma) ** 2)",
        ],
        "linear": [
            "def model(x, a, b):",
            '    """Linear: a * x + b"""',
            "    return a * x + b",
        ],
        "power_law": [
            "def model(x, a, b, c):",
            '    """Power law: a * x^b + c"""',
            "    return a * jnp.power(x, b) + c",
        ],
        "sigmoid": [
            "def model(x, L, k, x0, b):",
            '    """Sigmoid: L / (1 + exp(-k * (x - x0))) + b"""',
            "    return L / (1 + jnp.exp(-k * (x - x0))) + b",
        ],
    }

    if model_name in models:
        return models[model_name]
    else:
        return [
            f"# Built-in model: {model_name}",
            "from nlsq.models import get_builtin_model",
            f'model = get_builtin_model("{model_name}")',
        ]


def _get_polynomial_model_code(degree: int) -> list[str]:
    """Get Python code for a polynomial model."""
    params = ", ".join([f"c{i}" for i in range(degree + 1)])

    lines = [
        f"def model(x, {params}):",
        f'    """Polynomial of degree {degree}"""',
    ]

    # Build polynomial expression
    terms = []
    for i in range(degree + 1):
        if i == 0:
            terms.append("c0")
        elif i == 1:
            terms.append("c1 * x")
        else:
            terms.append(f"c{i} * x**{i}")

    expr = " + ".join(terms)
    lines.append(f"    return {expr}")

    return lines
