"""Objective wrappers and differential-evolution helpers for parameter
fitting."""

from __future__ import annotations
from datetime import datetime
from abc import ABC, abstractmethod
import csv
import sys
from pathlib import Path

import numpy as np
from scipy.optimize import differential_evolution


class ObjectiveWrapperBase(ABC):
    """Abstract objective wrapper with scaled optimization inputs."""

    def __init__(
        self,
        base_parameters: dict,
        keys_to_optimize: list[str],
        bounds_in_optimization: list[tuple[float, float]] | None = None,
        verbosity: bool = False,
    ) -> None:
        self.base_parameters = base_parameters
        self.keys_to_optimize = keys_to_optimize
        self.verbosity = verbosity
        self.bounds_in_optimization: list[tuple[float, float]] | None = None
        self.lower_bounds: np.ndarray | None = None
        self.upper_bounds: np.ndarray | None = None
        self.bound_ranges: np.ndarray | None = None
        if bounds_in_optimization is not None:
            self.set_bounds(bounds_in_optimization)

    def __call__(self, opt_values, verbosity: bool | None = None) -> float:
        """Evaluate the objective for scaled inputs."""
        current_args = self.base_parameters.copy()

        # opt_values are scaled (0: lower bound, 1: upper bound)
        opt_values_real = self.unscale(opt_values)

        for k, v in zip(self.keys_to_optimize, opt_values_real):
            current_args[k] = v

        if verbosity is None:
            verbosity = self.verbosity
        return self.objective(current_args, verbosity=verbosity)

    def set_bounds(self, bounds_in_optimization: list[tuple[float, float]]):
        """Set optimization bounds for scaling/unscaling."""
        self.bounds_in_optimization = bounds_in_optimization
        self.lower_bounds = np.array([b[0] for b in bounds_in_optimization])
        self.upper_bounds = np.array([b[1] for b in bounds_in_optimization])
        self.bound_ranges = self.upper_bounds - self.lower_bounds

    def scale(self, values_real: np.ndarray) -> np.ndarray:
        """Scale real values to unit interval using current bounds."""
        if self.lower_bounds is None or self.bound_ranges is None:
            raise ValueError("Bounds must be set before scaling values.")
        return (values_real - self.lower_bounds) / self.bound_ranges

    def unscale(self, values_scaled: np.ndarray) -> np.ndarray:
        """Unscale unit-interval values back to real units."""
        if self.lower_bounds is None or self.bound_ranges is None:
            raise ValueError("Bounds must be set before unscaling values.")
        return self.lower_bounds + self.bound_ranges * values_scaled

    @abstractmethod
    def objective(self, current_args: dict, verbosity: bool = False) -> float:
        """Compute objective value for the provided parameter dictionary."""
        raise NotImplementedError


class ParameterOptimization:
    """Optimization runner that requires an objective wrapper."""

    def __init__(self, wrapper: ObjectiveWrapperBase) -> None:
        self.wrapper = wrapper

    def build_initial_population(
        self,
        initial_guess_scaled: np.ndarray,
        popsize: int = 25,
        seed: int = 42,
        sigma: float = 0.10,
        sigma_random: float = 0.05,
    ) -> np.ndarray:
        """Build an initial population for differential evolution."""
        if self.wrapper.bounds_in_optimization is None:
            raise ValueError("Bounds must be set before building population.")
        n_params = len(self.wrapper.bounds_in_optimization)
        n_individuals = popsize * n_params

        population = np.empty((n_individuals, n_params), dtype=float)
        population[0] = initial_guess_scaled

        rng_local = np.random.default_rng(seed)
        half = (n_individuals - 1) // 2
        population[1 : 1 + half] = np.clip(
            initial_guess_scaled
            + rng_local.normal(sigma, sigma_random, size=(half, n_params)),
            0.0,
            1.0,
        )
        rest = n_individuals - 1 - half
        population[1 + half :] = rng_local.random((rest, n_params))
        return population

    # pylint: disable=too-many-arguments,too-many-positional-arguments
    def run_de(
        self,
        initial_guess_scaled: np.ndarray,
        popsize: int = 25,
        maxiter: int = 100,
        seed: int = 42,
        disp: bool | None = None,
        population: np.ndarray | None = None,
        log_generations: bool = True,
        prefix_file: str = "",
    ):
        """Run differential evolution in scaled space."""
        if self.wrapper.bounds_in_optimization is None:
            raise ValueError("Bounds must be set before running optimization.")
        if disp is None:
            disp = bool(self.wrapper.verbosity)
        n_params = len(self.wrapper.bounds_in_optimization)
        if population is None:
            population = self.build_initial_population(
                initial_guess_scaled, popsize=popsize, seed=seed
            )

        if log_generations:
            script_stem = Path(sys.argv[0]).stem or "de_run"
            csv_path = Path(f"{prefix_file}{script_stem}_de_gen.csv")
            header = [
                "generation",
                "timestamp",
                "fitness",
            ] + self.wrapper.keys_to_optimize
            if not csv_path.exists():
                with open(csv_path, "w", encoding="utf-8", newline="") as f:
                    csv.writer(f).writerow(header)

            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
            fitness0 = float(self.wrapper(initial_guess_scaled))
            params0 = self.wrapper.unscale(np.array(initial_guess_scaled))

            with open(csv_path, "a", encoding="utf-8", newline="") as f:
                csv.writer(f).writerow(
                    ["0-s", timestamp, fitness0]
                    + [float(v) for v in initial_guess_scaled]
                )
                csv.writer(f).writerow(
                    [0, timestamp, fitness0] + [float(v) for v in params0]
                )

            generation = 0

            def callback(xk, _convergence):
                nonlocal generation
                generation += 1
                timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
                fitness = float(self.wrapper(xk))
                params_unscaled = self.wrapper.unscale(np.array(xk))
                row = [generation, timestamp, fitness] + [
                    float(v) for v in params_unscaled
                ]
                row_scaled = [str(generation) + "-s", timestamp, fitness] + [
                    float(v) for v in xk
                ]
                with open(csv_path, "a", encoding="utf-8", newline="") as f:
                    csv.writer(f).writerow(row_scaled)
                    csv.writer(f).writerow(row)
                return False

            result = differential_evolution(
                self.wrapper,
                bounds=[(0.0, 1.0)] * n_params,
                workers=-1,
                callback=callback,
                polish=False,
                maxiter=maxiter,
            )

        else:
            result = differential_evolution(
                self.wrapper,
                bounds=[(0.0, 1.0)] * n_params,
                workers=-1,
                callback=None,
                polish=False,
                maxiter=maxiter,
            )

        if log_generations:
            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
            best_scaled = result.x
            best_unscaled = self.wrapper.unscale(best_scaled)
            best_fitness = float(result.fun)

            row = ["final", timestamp, best_fitness] + [
                float(v) for v in best_unscaled
            ]
            row_scaled = ["final-s", timestamp, best_fitness] + [
                float(v) for v in best_scaled
            ]

            with open(csv_path, "a", encoding="utf-8", newline="") as f:
                writer = csv.writer(f)
                writer.writerow(row_scaled)
                writer.writerow(row)
                writer.writerow(
                    ["status", "", result.message, f"nit={result.nit}"]
                )

        return result, population

    # pylint: disable=too-many-arguments,too-many-positional-arguments
    # pylint: disable=too-many-locals
    def run_optimization(
        self,
        initial_guess: np.ndarray,
        max_change_of_parameters: np.ndarray,
        *,
        verbosity: bool = False,
        de_kwargs: dict | None = None,
        population_kwargs: dict | None = None,
        top_n: int = 4,
    ):
        """Run differential evolution optimization."""
        de_kwargs = de_kwargs or {}
        population_kwargs = population_kwargs or {}

        bounds_in_optimization = bounds_from_center(
            initial_guess, max_change_of_parameters
        )
        self.wrapper.set_bounds(bounds_in_optimization)
        self.wrapper.verbosity = verbosity

        initial_guess_scaled = self.wrapper.scale(initial_guess)

        if verbosity:
            initial_args_scaled = dict(
                zip(self.wrapper.keys_to_optimize, initial_guess_scaled)
            )
            initial_args_unscaled = dict(
                zip(self.wrapper.keys_to_optimize, initial_guess)
            )
            print("Initial guess (scaled):")
            print(initial_args_scaled)
            print("Initial guess (unscaled):")
            print(initial_args_unscaled)
            print(
                "Initial score:",
                self.wrapper(initial_guess_scaled, verbosity=verbosity),
            )
            print("Start DE algorithm")

        population = self.build_initial_population(
            initial_guess_scaled, **population_kwargs
        )

        de_kwargs.setdefault("disp", verbosity)
        result, _ = self.run_de(
            initial_guess_scaled,
            population=population,
            **de_kwargs,
        )

        final_population = result.population
        fitness = result.population_energies
        sorted_indices = np.argsort(fitness)
        best_n = final_population[sorted_indices[:top_n]]
        best_n_f = fitness[sorted_indices[:top_n]]

        if verbosity:
            for idx_best, best in enumerate(best_n):
                optimized_args_scaled = dict(
                    zip(self.wrapper.keys_to_optimize, best)
                )
                optimized_args_unscaled = dict(
                    zip(
                        self.wrapper.keys_to_optimize,
                        self.wrapper.unscale(best),
                    )
                )
                print(f"Optimized parameters best {idx_best + 1}:")
                print(optimized_args_scaled)
                print(optimized_args_unscaled)
                print("Final score:", best_n_f[idx_best])

        return {
            "de": result,
            "best_scaled": best_n,
            "best_scores": best_n_f,
            "best_unscaled": np.array(
                [self.wrapper.unscale(b) for b in best_n]
            ),
        }


def bounds_from_center(center, deltas):
    """Return symmetric bounds around a center vector."""
    return [(val - delta, val + delta) for val, delta in zip(center, deltas)]
