# -*- coding: utf-8 -*-
"""Optimization helpers for recompression Brayton cycle studies."""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize_scalar

from .cycles import BraytonCycle
from .sensitivity import RC_FIXED_KEYS, RC_PARAM_KEYS


INVALID_OBJECTIVE = 1.0e12


def _require_keys(data, required_keys, data_name):
    missing = [key for key in required_keys if key not in data]
    if missing:
        missing_text = ", ".join(missing)
        raise ValueError(f"{data_name} missing required keys: {missing_text}")


def _resolve_variable_location(variable_name, fixed_params, params):
    in_fixed = variable_name in fixed_params
    in_params = variable_name in params

    if in_fixed and in_params:
        raise ValueError(
            f"{variable_name!r} exists in both fixed_params and params; "
            "rename one of them or choose the target explicitly."
        )
    if in_fixed:
        return "fixed"
    if in_params:
        return "params"
    raise ValueError(f"Unknown variable: {variable_name}")


def _with_updated_value(fixed_params, params, variable_name, value):
    fixed = dict(fixed_params)
    cycle_params = dict(params)
    location = _resolve_variable_location(variable_name, fixed, cycle_params)

    if location == "fixed":
        fixed[variable_name] = value
    else:
        cycle_params[variable_name] = value

    return fixed, cycle_params


def _validate_bounds(bounds, bounds_name):
    if len(bounds) != 2:
        raise ValueError(f"{bounds_name} must contain exactly two values")
    if bounds[0] >= bounds[1]:
        raise ValueError(f"{bounds_name} lower bound must be smaller than upper bound")


def _validate_num_points(num_points):
    if num_points < 1:
        raise ValueError("num_points must be at least 1")


def _evaluate_rc_cycle(fixed_params, params, refprop_path=None):
    fixed = dict(fixed_params)
    cycle_params = dict(params)

    _require_keys(fixed, RC_FIXED_KEYS, "fixed_params")
    _require_keys(cycle_params, RC_PARAM_KEYS, "params")

    cycle_kwargs = {"name": "rc optimization evaluation"}
    if refprop_path is not None:
        cycle_kwargs["refprop_path"] = refprop_path

    cycle = BraytonCycle(**cycle_kwargs)
    efficiency = cycle.RC(
        T_low=fixed["T_low"],
        T_high=fixed["T_high"],
        p_low=fixed["p_low"],
        p_high=fixed["p_high"],
        ploss=fixed["ploss"],
        param=cycle_params,
    )
    return efficiency, cycle


def _pinch_points_ok(cycle):
    if not cycle.recuperator:
        return True
    return all(
        recuperator.check_pinch_point(cycle.property_calculator)
        for recuperator in cycle.recuperator
    )


def _rc_objective(fixed_params, params, refprop_path=None, check_pinch=False):
    try:
        efficiency, cycle = _evaluate_rc_cycle(fixed_params, params, refprop_path)
        if check_pinch and not _pinch_points_ok(cycle):
            return INVALID_OBJECTIVE
        return -efficiency
    except Exception:
        return INVALID_OBJECTIVE


def _annotate_result(result):
    result.valid = bool(
        result.success
        and np.isfinite(result.fun)
        and result.fun < INVALID_OBJECTIVE / 2
    )
    result.best_efficiency = -result.fun if result.valid else np.nan
    return result


def optimize_rc_param(
    fixed_params,
    params,
    target_var_name,
    bounds,
    refprop_path=None,
    check_pinch=False,
):
    """Optimize one RC component/cycle parameter for maximum efficiency."""
    if target_var_name not in params:
        raise ValueError(f"{target_var_name!r} is not in params")
    _validate_bounds(bounds, "bounds")

    def objective(value):
        trial_params = dict(params)
        trial_params[target_var_name] = value
        return _rc_objective(
            fixed_params,
            trial_params,
            refprop_path=refprop_path,
            check_pinch=check_pinch,
        )

    result = minimize_scalar(objective, bounds=bounds, method="bounded")
    return _annotate_result(result)


def optimize_rc_fixed_param(
    fixed_params,
    params,
    target_var_name,
    bounds,
    refprop_path=None,
    check_pinch=False,
):
    """Optimize one RC boundary-condition parameter for maximum efficiency."""
    if target_var_name not in fixed_params:
        raise ValueError(f"{target_var_name!r} is not in fixed_params")
    _validate_bounds(bounds, "bounds")

    def objective(value):
        trial_fixed = dict(fixed_params)
        trial_fixed[target_var_name] = value
        return _rc_objective(
            trial_fixed,
            params,
            refprop_path=refprop_path,
            check_pinch=check_pinch,
        )

    result = minimize_scalar(objective, bounds=bounds, method="bounded")
    return _annotate_result(result)


def scan_rc_efficiency(
    fixed_params,
    params,
    target_var_name,
    bounds,
    num_points=50,
    refprop_path=None,
):
    """Evaluate RC efficiency over a one-dimensional variable sweep."""
    _validate_bounds(bounds, "bounds")
    _validate_num_points(num_points)

    x_values = []
    efficiencies = []
    errors = []

    for value in np.linspace(bounds[0], bounds[1], num_points):
        trial_fixed, trial_params = _with_updated_value(
            fixed_params, params, target_var_name, value
        )
        try:
            efficiency, _ = _evaluate_rc_cycle(
                trial_fixed, trial_params, refprop_path=refprop_path
            )
        except Exception as exc:
            errors.append((value, exc))
            continue

        x_values.append(value)
        efficiencies.append(efficiency)

    return x_values, efficiencies, errors


def sweep_and_optimize_rc(
    fixed_params,
    params,
    sweep_var,
    sweep_bounds,
    opt_var,
    opt_bounds,
    num_points=50,
    refprop_path=None,
    check_pinch=True,
):
    """Sweep one variable and optimize another at each sweep point."""
    if sweep_var == opt_var:
        raise ValueError("sweep_var and opt_var must be different variables")
    _validate_bounds(sweep_bounds, "sweep_bounds")
    _validate_bounds(opt_bounds, "opt_bounds")
    _validate_num_points(num_points)

    valid_sweep_vals = []
    best_efficiencies = []
    best_opt_vals = []
    failures = []

    for sweep_value in np.linspace(sweep_bounds[0], sweep_bounds[1], num_points):
        current_fixed, current_params = _with_updated_value(
            fixed_params, params, sweep_var, sweep_value
        )

        def objective(opt_value):
            trial_fixed, trial_params = _with_updated_value(
                current_fixed, current_params, opt_var, opt_value
            )
            return _rc_objective(
                trial_fixed,
                trial_params,
                refprop_path=refprop_path,
                check_pinch=check_pinch,
            )

        result = minimize_scalar(objective, bounds=opt_bounds, method="bounded")
        result = _annotate_result(result)

        if result.valid:
            valid_sweep_vals.append(sweep_value)
            best_efficiencies.append(result.best_efficiency)
            best_opt_vals.append(result.x)
        else:
            failures.append((sweep_value, result))

    return valid_sweep_vals, best_efficiencies, best_opt_vals, failures


def plot_optimization_landscape(
    fixed_params,
    params,
    target_var_name,
    bounds,
    result=None,
    num_points=50,
    refprop_path=None,
    show=True,
):
    """Plot RC efficiency over a one-dimensional sweep."""
    x_values, efficiencies, errors = scan_rc_efficiency(
        fixed_params,
        params,
        target_var_name,
        bounds,
        num_points=num_points,
        refprop_path=refprop_path,
    )
    efficiency_percent = [efficiency * 100 for efficiency in efficiencies]

    fig, ax = plt.subplots(figsize=(8, 6), dpi=120)
    ax.plot(x_values, efficiency_percent, color="#1f77b4", linewidth=2)
    ax.set_title(f"RC efficiency vs {target_var_name}")
    ax.set_xlabel(target_var_name)
    ax.set_ylabel("Cycle efficiency (%)")
    ax.grid(True, linestyle=":", alpha=0.7)

    if result is not None and getattr(result, "valid", result.success):
        best_efficiency = getattr(result, "best_efficiency", -result.fun)
        ax.scatter(
            result.x,
            best_efficiency * 100,
            color="red",
            marker="*",
            s=200,
            zorder=5,
        )
        ax.axvline(x=result.x, color="gray", linestyle="--", alpha=0.6)
        ax.axhline(y=best_efficiency * 100, color="gray", linestyle="--", alpha=0.6)

    fig.tight_layout()
    if show:
        plt.show()

    return fig, ax, x_values, efficiency_percent, errors


def plot_sweep_optimization_results(
    sweep_var,
    sweep_values,
    best_efficiencies,
    opt_var,
    best_opt_values,
    show=True,
):
    """Plot nested sweep and optimization results with two y axes."""
    if not sweep_values:
        raise ValueError("No valid sweep data to plot")

    fig, ax1 = plt.subplots(figsize=(9, 6), dpi=120)

    color1 = "#1f77b4"
    ax1.set_xlabel(sweep_var)
    ax1.set_ylabel("Best cycle efficiency (%)", color=color1)
    best_efficiency_percent = [
        efficiency * 100 for efficiency in best_efficiencies
    ]
    ax1.plot(sweep_values, best_efficiency_percent, color=color1, linewidth=2.5)
    ax1.tick_params(axis="y", labelcolor=color1)
    ax1.grid(True, linestyle=":", alpha=0.6)

    ax2 = ax1.twinx()
    color2 = "#d62728"
    ax2.set_ylabel(f"Best {opt_var}", color=color2)
    ax2.plot(
        sweep_values,
        best_opt_values,
        color=color2,
        linestyle="--",
        linewidth=2,
    )
    ax2.tick_params(axis="y", labelcolor=color2)

    fig.tight_layout()
    if show:
        plt.show()

    return fig, ax1, ax2