import json
import random
from pathlib import Path
from typing import List, Tuple

from simulation_model import SimulationModel

# Decision variables: two factory factors only
TARGET_FACTORIES = ["ZhongcaiBaicheng", "ZhongcaiHami"]

# Bounds
FACTOR_BOUNDS = (0.8, 3.0)

POP_SIZE = 20
GENERATIONS = 50
MUTATION_RATE = 0.2
MUTATION_STD = 0.05  # for factors


def clip(val: float, bounds: Tuple[float, float]) -> float:
    lo, hi = bounds
    return max(lo, min(hi, val))


def evaluate(genes: List[float]) -> float:
    factory_factors = {fid: val for fid, val in zip(TARGET_FACTORIES, genes)}
    model = SimulationModel(factory_factors=factory_factors, output_enabled=False)
    while model.running:
        model.step()
    return model.mean_abs_error


def mutate(genes: List[float]) -> List[float]:
    new = genes.copy()
    for i in range(len(new)):
        if random.random() < MUTATION_RATE:
            jitter = random.gauss(0, MUTATION_STD)
            new[i] = clip(new[i] + jitter, FACTOR_BOUNDS)
    return new


def crossover(p1: List[float], p2: List[float]) -> Tuple[List[float], List[float]]:
    point = random.randint(1, len(p1) - 1)
    c1 = p1[:point] + p2[point:]
    c2 = p2[:point] + p1[point:]
    return c1, c2


def init_population() -> List[List[float]]:
    pop = []
    best_path = Path("data") / "ga_two_best_params.json"
    seed_indiv = None
    if best_path.exists():
        try:
            best = json.loads(best_path.read_text(encoding="utf-8"))
            seed_indiv = [float(best.get(f"factor_{fid}", random.uniform(*FACTOR_BOUNDS))) for fid in TARGET_FACTORIES]
        except Exception:
            seed_indiv = None

    for idx in range(POP_SIZE):
        if seed_indiv is not None and idx == 0:
            pop.append(seed_indiv)
            continue
        indiv = [random.uniform(*FACTOR_BOUNDS) for _ in TARGET_FACTORIES]
        pop.append(indiv)
    return pop


def main():
    best_genes = None
    best_score = float("inf")
    population = init_population()

    for gen in range(GENERATIONS):
        scored = []
        for indiv in population:
            score = evaluate(indiv)
            scored.append((score, indiv))
            if score < best_score:
                best_score = score
                best_genes = indiv
        scored.sort(key=lambda x: x[0])
        next_pop = [scored[0][1], scored[1][1]]
        while len(next_pop) < POP_SIZE:
            parents = random.sample(scored[:max(3, len(scored))], 2)
            c1, c2 = crossover(parents[0][1], parents[1][1])
            next_pop.append(mutate(c1))
            if len(next_pop) < POP_SIZE:
                next_pop.append(mutate(c2))
        population = next_pop
        print(f"Generation {gen+1}: best={best_score:.4f}")

    result = {}
    for fid, val in zip(TARGET_FACTORIES, best_genes):
        result[f"factor_{fid}"] = val

    data_path = Path("data") / "ga_two_best_params.json"
    prev_score = float("inf")
    if data_path.exists():
        try:
            prev = json.loads(data_path.read_text(encoding="utf-8"))
            prev_score = float(prev.get("best_score", float("inf")))
        except Exception:
            prev_score = float("inf")

    if best_score < prev_score:
        data_path.write_text(
            json.dumps({"best_score": best_score, **result}, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        print(f"New best mean abs error: {best_score:.4f}, saved to {data_path}")
    else:
        print(f"Best mean abs error: {best_score:.4f} (not better than {prev_score:.4f}, not saved)")


if __name__ == "__main__":
    main()