"""Baseline optimizers backed by the mealpy library (v3).

Each class preserves the original interface::

    best_pos, best_fit, convergence = optimizer.optimize(
        objective_fn, bounds, dim, max_iterations
    )

All algorithms are sourced from the ``mealpy`` package
(https://github.com/thieu1995/mealpy).
"""

import numpy as np
from mealpy import (
    FloatVar,
    PSO  as _PSO,
    GWO  as _GWO,
    GA   as _GA,
    DE   as _DE,
    SA   as _SA,
    WOA  as _WOA,
    ABC  as _ABC,
    FFA  as _FFA,   # Firefly Algorithm (mealpy abbreviates it FFA)
    ACOR as _ACOR,  # ACO for Continuous domains
)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _bounds_arrays(bounds, dim):
    """Return (lb, ub) as float64 numpy arrays of length *dim*."""
    lb = np.array(bounds[0]) if hasattr(bounds[0], '__len__') else np.full(dim, bounds[0])
    ub = np.array(bounds[1]) if hasattr(bounds[1], '__len__') else np.full(dim, bounds[1])
    return lb.astype(float), ub.astype(float)


def _make_problem(obj_fn, lb, ub):
    """Build a mealpy minimisation problem dict."""
    return {
        "obj_func": obj_fn,
        "bounds":   FloatVar(lb=lb.tolist(), ub=ub.tolist()),
        "minmax":   "min",
        "log_to":   None,
        "verbose":  False,
    }


def _run(model, obj_fn, lb, ub):
    """Solve *problem* and return ``(best_pos, best_fit, convergence)``."""
    result     = model.solve(_make_problem(obj_fn, lb, ub))
    best_pos   = np.array(result.solution)
    best_fit   = float(result.target.fitness)
    convergence = list(model.history.list_global_best_fit)
    return best_pos, best_fit, convergence


# ---------------------------------------------------------------------------
# Baseline classes
# ---------------------------------------------------------------------------

class PSO:
    """Kennedy & Eberhart (1995) – Particle Swarm Optimization.

    Backend: ``mealpy.PSO.OriginalPSO``
    """

    def __init__(self, pop_size=30, c1=2.0, c2=2.0, w=0.7):
        self.pop_size = pop_size
        self.c1, self.c2, self.w = c1, c2, w

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _PSO.OriginalPSO(
            epoch=max_iterations, pop_size=self.pop_size,
            c1=self.c1, c2=self.c2, w=self.w,
        )
        return _run(model, objective_fn, lb, ub)


class GWO:
    """Mirjalili et al. (2014) – Grey Wolf Optimizer.

    Backend: ``mealpy.GWO.OriginalGWO``
    """

    def __init__(self, pop_size=30):
        self.pop_size = pop_size

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _GWO.OriginalGWO(epoch=max_iterations, pop_size=self.pop_size)
        return _run(model, objective_fn, lb, ub)


class GA:
    """Holland (1975) – Genetic Algorithm.

    Backend: ``mealpy.GA.BaseGA``
    """

    def __init__(self, pop_size=30, mutation_rate=0.025, crossover_rate=0.95):
        self.pop_size = pop_size
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _GA.BaseGA(
            epoch=max_iterations, pop_size=self.pop_size,
            pc=self.crossover_rate, pm=self.mutation_rate,
        )
        return _run(model, objective_fn, lb, ub)


class DE:
    """Storn & Price (1997) – Differential Evolution.

    Backend: ``mealpy.DE.OriginalDE``
    """

    def __init__(self, pop_size=30, F=0.8, CR=0.9):
        self.pop_size = pop_size
        self.F  = F   # differential weight
        self.CR = CR  # crossover probability

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _DE.OriginalDE(
            epoch=max_iterations, pop_size=self.pop_size,
            wf=self.F, cr=self.CR,
        )
        return _run(model, objective_fn, lb, ub)


class SA:
    """Kirkpatrick et al. (1983) – Simulated Annealing.

    Backend: ``mealpy.SA.OriginalSA``
    """

    def __init__(self, temp_start=100.0, temp_min=1e-3, cooling=0.98, step_scale=0.1):
        self.temp_start = temp_start
        self.temp_min   = temp_min
        self.cooling    = cooling
        self.step_scale = step_scale

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _SA.OriginalSA(
            epoch=max_iterations, pop_size=2,
            temp_init=self.temp_start, step_size=self.step_scale,
        )
        return _run(model, objective_fn, lb, ub)


class JADE:
    """Zhang & Sanderson (2009) – Adaptive Differential Evolution (JADE).

    Backend: ``mealpy.DE.JADE``
    """

    def __init__(self, pop_size=30, p_best=0.1, c=0.1):
        self.pop_size = pop_size
        self.p_best   = p_best
        self.c        = c

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _DE.JADE(
            epoch=max_iterations, pop_size=self.pop_size,
            pt=self.p_best, ap=self.c,
        )
        return _run(model, objective_fn, lb, ub)


class WOA:
    """Mirjalili & Lewis (2016) – Whale Optimization Algorithm.

    Backend: ``mealpy.WOA.OriginalWOA``
    """

    def __init__(self, pop_size=30):
        self.pop_size = pop_size

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _WOA.OriginalWOA(epoch=max_iterations, pop_size=self.pop_size)
        return _run(model, objective_fn, lb, ub)


class ABC:
    """Karaboga (2005) – Artificial Bee Colony.

    Backend: ``mealpy.ABC.OriginalABC``
    """

    def __init__(self, pop_size=30, limit=50):
        self.pop_size = pop_size
        self.limit    = limit

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        # mealpy n_limits is per-source; scale from absolute limit
        n_limits = max(1, self.limit // max(1, self.pop_size))
        model = _ABC.OriginalABC(
            epoch=max_iterations, pop_size=self.pop_size,
            n_limits=n_limits,
        )
        return _run(model, objective_fn, lb, ub)


class FA:
    """Yang (2008) – Firefly Algorithm.

    Backend: ``mealpy.FFA.OriginalFFA``
    """

    def __init__(self, pop_size=30, alpha=0.2, beta=1.0, gamma=0.001):
        self.pop_size = pop_size
        self.alpha = alpha  # randomisation parameter
        self.beta  = beta   # base attractiveness
        self.gamma = gamma  # light absorption coefficient

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _FFA.OriginalFFA(
            epoch=max_iterations, pop_size=self.pop_size,
            gamma=self.gamma, beta_base=self.beta, alpha=self.alpha,
        )
        return _run(model, objective_fn, lb, ub)


class ACO:
    """Socha & Dorigo (2008) – Ant Colony Optimisation for Continuous Domains (ACO-R).

    Backend: ``mealpy.ACOR.OriginalACOR``
    """

    def __init__(self, pop_size=30, sample_count=25, intent_factor=0.5, zeta=1.0):
        self.pop_size      = pop_size
        self.sample_count  = sample_count
        self.intent_factor = intent_factor
        self.zeta          = zeta

    def optimize(self, objective_fn, bounds, dim, max_iterations=200):
        lb, ub = _bounds_arrays(bounds, dim)
        model = _ACOR.OriginalACOR(
            epoch=max_iterations, pop_size=self.pop_size,
            sample_count=self.sample_count,
            intent_factor=self.intent_factor,
            zeta=self.zeta,
        )
        return _run(model, objective_fn, lb, ub)