"""
Optimization utilities for fitting Slater-Koster parameters.

This is the core module for inverse problems: fitting TB parameters
to match experimental data (ARPES, band structure, etc.).
"""

from __future__ import annotations

import time

import jax
import jax.numpy as jnp
import optax
from beartype.typing import Callable, Dict, NamedTuple, Optional, Tuple
from jax import Array


class OptimizationResult(NamedTuple):
    """
    Result of parameter optimization.

    Attributes
    ----------
    parameters : Dict[str, Array]
        Optimized parameters.
    loss_history : Array
        Loss value at each iteration.
    grad_history : Array
        Gradient norm at each iteration.
    n_iterations : int
        Number of iterations performed.
    converged : bool
        Whether optimization converged.
    final_loss : float
        Final loss value.
    time_elapsed : float
        Total time in seconds.
    """

    parameters: Dict[str, Array]
    loss_history: Array
    grad_history: Array
    n_iterations: int
    converged: bool
    final_loss: float
    time_elapsed: float


class ParameterSpec(NamedTuple):
    """
    Specification for a fittable parameter.

    Attributes
    ----------
    name : str
        Parameter name.
    initial_value : float
        Starting value.
    bounds : Tuple[float, float], optional
        (min, max) bounds or None for unbounded.
    fixed : bool
        If True, parameter is not optimized.
    transform : str, optional
        Optional transform ('log', 'sigmoid', etc.).
    """

    name: str
    initial_value: float
    bounds: Optional[Tuple[float, float]] = None
    fixed: bool = False
    transform: Optional[str] = None


# =============================================================================
# Main fitting function
# =============================================================================


def fit_parameters(
    loss_fn: Callable[[Dict[str, Array]], float],
    initial_params: Dict[str, Array],
    param_specs: Optional[Dict[str, ParameterSpec]] = None,
    optimizer: str = "adam",
    learning_rate: float = 0.01,
    max_iterations: int = 1000,
    tolerance: float = 1e-6,
    patience: int = 50,
    verbose: bool = True,
    callback: Optional[Callable] = None,
) -> OptimizationResult:
    """
    Fit parameters by minimizing a loss function.

    This is the main interface for inverse problems in DiffTB.

    Parameters
    ----------
    loss_fn : Callable[[Dict[str, Array]], float]
        Function params -> scalar loss (must be JAX-differentiable).
    initial_params : Dict[str, Array]
        Dictionary of initial parameter values.
    param_specs : Dict[str, ParameterSpec], optional
        Optional specifications for each parameter.
    optimizer : str, optional
        Optimizer name ('adam', 'sgd', 'adamw'). Default is 'adam'.
    learning_rate : float, optional
        Learning rate. Default is 0.01.
    max_iterations : int, optional
        Maximum iterations. Default is 1000.
    tolerance : float, optional
        Convergence tolerance for loss change. Default is 1e-6.
    patience : int, optional
        Stop if no improvement for this many iterations. Default is 50.
    verbose : bool, optional
        Print progress. Default is True.
    callback : Callable, optional
        Optional function called each iteration with (iter, params, loss).

    Returns
    -------
    result : OptimizationResult
        OptimizationResult with optimized parameters and history.
    """
    start_time = time.time()

    # Separate fixed and trainable parameters
    trainable_params = {}
    fixed_params = {}

    for name, value in initial_params.items():
        if param_specs and name in param_specs and param_specs[name].fixed:
            fixed_params[name] = value
        else:
            trainable_params[name] = value

    # Create loss function that includes fixed params
    def trainable_loss(trainable_params):
        all_params = {**trainable_params, **fixed_params}
        return loss_fn(all_params)

    # Setup optimizer
    if optimizer == "adam":
        opt = optax.adam(learning_rate)
    elif optimizer == "sgd":
        opt = optax.sgd(learning_rate, momentum=0.9)
    elif optimizer == "adamw":
        opt = optax.adamw(learning_rate)
    else:
        raise ValueError(f"Unknown optimizer: {optimizer}")

    opt_state = opt.init(trainable_params)

    # JIT compile the update step
    @jax.jit
    def update_step(params, opt_state):
        loss, grads = jax.value_and_grad(trainable_loss)(params)
        updates, opt_state = opt.update(grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        grad_norm = jnp.sqrt(
            sum(jnp.sum(g**2) for g in jax.tree_util.tree_leaves(grads))
        )
        return params, opt_state, loss, grad_norm

    # Optimization loop
    loss_history = []
    grad_history = []
    best_loss = float("inf")
    best_params = trainable_params
    no_improvement = 0
    converged = False

    for iteration in range(max_iterations):
        trainable_params, opt_state, loss, grad_norm = update_step(
            trainable_params, opt_state
        )

        loss_val = float(loss)
        loss_history.append(loss_val)
        grad_history.append(float(grad_norm))

        # Track best
        if loss_val < best_loss - tolerance:
            best_loss = loss_val
            best_params = jax.tree_util.tree_map(
                lambda x: x.copy(), trainable_params
            )
            no_improvement = 0
        else:
            no_improvement += 1

        # Logging
        if verbose and iteration % 100 == 0:
            print(
                f"Iter {iteration:5d}: loss = {loss_val:.6f}, |grad| = {grad_norm:.6f}"
            )

        # Callback
        if callback is not None:
            callback(iteration, {**trainable_params, **fixed_params}, loss_val)

        # Check convergence
        if grad_norm < tolerance:
            converged = True
            if verbose:
                print(f"Converged: gradient norm < {tolerance}")
            break

        if no_improvement >= patience:
            if verbose:
                print(
                    f"Early stopping: no improvement for {patience} iterations"
                )
            break

    time_elapsed = time.time() - start_time

    # Combine best trainable with fixed
    final_params = {**best_params, **fixed_params}

    return OptimizationResult(
        parameters=final_params,
        loss_history=jnp.array(loss_history),
        grad_history=jnp.array(grad_history),
        n_iterations=iteration + 1,
        converged=converged,
        final_loss=best_loss,
        time_elapsed=time_elapsed,
    )


# =============================================================================
# Hierarchical fitting
# =============================================================================


def hierarchical_fit(
    loss_fn: Callable,
    param_groups: list[Dict[str, Array]],
    group_names: list[str],
    **fit_kwargs,
) -> Tuple[Dict[str, Array], list[OptimizationResult]]:
    """
    Fit parameters in stages, freezing earlier groups.

    This is the key approach for transferability testing:

    1. Fit Bi-Te parameters on Bi2Te3
    2. Freeze those, fit only Mn parameters on MnBi2Te4

    Parameters
    ----------
    loss_fn : Callable
        Loss function taking all parameters.
    param_groups : list[Dict[str, Array]]
        List of parameter dictionaries for each stage.
    group_names : list[str]
        Names for logging.
    **fit_kwargs
        Arguments passed to fit_parameters.

    Returns
    -------
    final_params : Dict[str, Array]
        All optimized parameters.
    results : list[OptimizationResult]
        OptimizationResult for each stage.
    """
    all_params = {}
    results = []

    for stage, (group, name) in enumerate(
        zip(param_groups, group_names, strict=False)
    ):
        print(f"\n{'='*50}")
        print(f"Stage {stage + 1}: Fitting {name}")
        print(f"{'='*50}")

        # Combine with previously fitted (frozen) parameters
        current_params = {**all_params, **group}

        # Create specs that freeze previously fitted params
        param_specs = {}
        for param_name, param_value in all_params.items():
            param_specs[param_name] = ParameterSpec(
                name=param_name,
                initial_value=float(param_value),
                fixed=True,
            )

        result = fit_parameters(
            loss_fn, current_params, param_specs=param_specs, **fit_kwargs
        )

        results.append(result)
        all_params = result.parameters

        print(f"Stage {stage + 1} complete: loss = {result.final_loss:.6f}")

    return all_params, results


# =============================================================================
# Sensitivity analysis
# =============================================================================


@jax.jit
def compute_parameter_sensitivity(
    loss_fn: Callable, params: Dict[str, Array]
) -> Dict[str, Array]:
    """
    Compute gradient of loss with respect to each parameter.

    This reveals which parameters most strongly affect the loss,
    guiding which parameters to prioritize in fitting.

    Parameters
    ----------
    loss_fn : Callable
        Loss function.
    params : Dict[str, Array]
        Current parameter values.

    Returns
    -------
    gradients : Dict[str, Array]
        Dictionary of gradients for each parameter.
    """
    return jax.grad(loss_fn)(params)


def sensitivity_analysis(
    loss_fn: Callable, params: Dict[str, Array], perturbation: float = 0.01
) -> Dict[str, float]:
    """
    Analyze parameter sensitivity by finite differences.

    Computes |∂L/∂p| × |p| for each parameter (relative sensitivity).

    Parameters
    ----------
    loss_fn : Callable
        Loss function.
    params : Dict[str, Array]
        Parameter values.
    perturbation : float, optional
        Fractional perturbation size. Default is 0.01.

    Returns
    -------
    sensitivities : Dict[str, float]
        Dictionary of relative sensitivities.
    """
    loss_fn(params)
    sensitivities = {}

    for name, value in params.items():
        # Perturb parameter
        delta = perturbation * jnp.abs(value) + 1e-8

        perturbed = params.copy()
        perturbed[name] = value + delta
        loss_plus = loss_fn(perturbed)

        perturbed[name] = value - delta
        loss_minus = loss_fn(perturbed)

        # Central difference
        dL_dp = (loss_plus - loss_minus) / (2 * delta)

        # Relative sensitivity
        sensitivities[name] = float(jnp.abs(dL_dp * value))

    return sensitivities


def parameter_hessian(loss_fn: Callable, params: Dict[str, Array]) -> Array:
    """
    Compute Hessian matrix of loss with respect to parameters.

    Useful for understanding parameter correlations and
    identifiability issues.

    Parameters
    ----------
    loss_fn : Callable
        Loss function.
    params : Dict[str, Array]
        Parameter values (must be flat array).

    Returns
    -------
    hessian : Array
        Hessian matrix.
    """
    # Convert dict to flat array for Hessian computation
    param_list = list(params.values())
    param_flat = jnp.concatenate([p.ravel() for p in param_list])

    def flat_loss(flat_params):
        # Reconstruct dict
        idx = 0
        reconstructed = {}
        for name, orig in params.items():
            size = orig.size
            reconstructed[name] = flat_params[idx : idx + size].reshape(
                orig.shape
            )
            idx += size
        return loss_fn(reconstructed)

    return jax.hessian(flat_loss)(param_flat)


# =============================================================================
# Cross-validation
# =============================================================================


def cross_validate(
    loss_fn_factory: Callable[[Array], Callable],
    data_splits: list[Tuple[Array, Array]],
    initial_params: Dict[str, Array],
    **fit_kwargs,
) -> Tuple[float, float, list[OptimizationResult]]:
    """
    K-fold cross-validation for parameter fitting.

    Parameters
    ----------
    loss_fn_factory : Callable[[Array], Callable]
        Function that creates loss function from training data.
    data_splits : list[Tuple[Array, Array]]
        List of (train_indices, val_indices) for each fold.
    initial_params : Dict[str, Array]
        Initial parameters.
    **fit_kwargs
        Arguments for fit_parameters.

    Returns
    -------
    mean_val_loss : float
        Mean validation loss across folds.
    std_val_loss : float
        Standard deviation of validation loss.
    results : list[OptimizationResult]
        Results from each fold.
    """
    val_losses = []
    results = []

    for fold, (train_idx, val_idx) in enumerate(data_splits):
        print(f"\nFold {fold + 1}/{len(data_splits)}")

        # Create loss function for this fold's training data
        train_loss_fn = loss_fn_factory(train_idx)
        val_loss_fn = loss_fn_factory(val_idx)

        # Fit on training data
        result = fit_parameters(
            train_loss_fn, initial_params.copy(), **fit_kwargs
        )
        results.append(result)

        # Evaluate on validation data
        val_loss = val_loss_fn(result.parameters)
        val_losses.append(float(val_loss))

        print(f"  Train loss: {result.final_loss:.6f}")
        print(f"  Val loss: {val_loss:.6f}")

    mean_val = jnp.mean(jnp.array(val_losses))
    std_val = jnp.std(jnp.array(val_losses))

    return float(mean_val), float(std_val), results
