"""
xp3: iterative surrogate refinement via Phase-2 schedule capture.

Pipeline (one outer loop per init config, one inner loop per iteration):
  iter 0 : build initial dataset per init config → Phase 1 → Phase 2
  iter k>0: pool = initial + all unique captures+swaps so far → fresh 90/10 resplit → Phase 1 → Phase 2

Initial dataset modes (set in INIT_CONFIGS):
  "worst"            - keep worst_pct% highest-makespan schedules
  "uniform"          - uniform random sample of pct% of all schedules
  "percentile_range" - keep schedules whose makespan falls between low_pct% and high_pct% percentile

Run from:  src/xp3/
"""

import os
import sys
import json
import shutil
import math

import numpy as np
from tqdm import tqdm

# ── resolve paths ─────────────────────────────────────────────────────────────
XP3_DIR  = os.path.dirname(os.path.abspath(__file__))
XP1_DIR  = os.path.abspath(os.path.join(XP3_DIR, "../xp1"))
sys.path.insert(0, XP1_DIR)   # for train, GPT, create_dataset
sys.path.insert(0, XP3_DIR)   # xp3's recover_schedules takes priority over xp1's

# viz notebooks (viz_train.ipynb, viz_rs.ipynb) are resolved relative to cwd
os.chdir(XP1_DIR)

from train import train
from recover_schedules import recover_schedules   # xp3's modified version
from create_dataset import evaluate_makespan      # cumulative makespan computation

RAW_DATA_DIR = os.path.join(XP1_DIR, "outputs/exhaustive_9_6")
XP3_OUT_ROOT = os.path.join(XP3_DIR, "outputs/exhaustive_9_6")

# ── experiment parameters ─────────────────────────────────────────────────────

# Each entry defines one initial-dataset configuration.
# Supported modes:
#   {"mode": "worst",            "pct": <float>}
#   {"mode": "uniform",          "pct": <float>}
#   {"mode": "percentile_range", "low_pct": <float>, "high_pct": <float>}
INIT_CONFIGS = [
    {"mode": "uniform",   "pct": 5},
    # {"mode": "worst",   "pct": 10},
    # {"mode": "worst",   "pct": 20},
    # {"mode": "worst",   "pct": 50},
]

N_ITERATIONS          = 10   # number of Phase1→Phase2 cycles per init config
MAX_RETRIES           = 3    # max Phase-2 relaunches per iteration if no new schedules found
N_SWAPS_PER_SCHEDULE  = 1000    # swap-augmented schedules generated per each newly captured schedule
N_EMBD, N_HEAD, N_LAYER = 40, 4, 2


# ── helpers ───────────────────────────────────────────────────────────────────

def _config_key(cfg):
    """Return a filesystem-safe string identifying an init config."""
    mode = cfg["mode"]
    if mode == "worst":
        return f"worst_pct_{cfg['pct']}"
    if mode == "uniform":
        return f"uniform_pct_{cfg['pct']}"
    if mode == "percentile_range":
        return f"pct_range_{cfg['low_pct']}_{cfg['high_pct']}"
    raise ValueError(f"Unknown mode: {mode}")


def _stratified_split(schedules, makespans, train_ratio, seed):
    """90/10 stratified split by final makespan value. Returns (train_idx, val_idx)."""
    ms = makespans[:, -1]
    np.random.seed(seed)
    train_idx_list, val_idx_list = [], []
    for m in np.unique(ms):
        group = np.where(ms == m)[0]
        np.random.shuffle(group)
        n_train = max(1, int(len(group) * train_ratio))
        if n_train >= len(group):
            train_idx_list.append(group)
        else:
            train_idx_list.append(group[:n_train])
            val_idx_list.append(group[n_train:])
    train_idx = np.concatenate(train_idx_list)
    val_idx   = np.concatenate(val_idx_list) if val_idx_list else np.array([], dtype=int)
    return train_idx, val_idx


def _save_split(out_dir, schedules, makespans, raw_data_dir, extra_meta=None):
    """Write a processed dataset directory from a (schedules, makespans) pool."""
    os.makedirs(out_dir, exist_ok=True)
    train_idx, val_idx = _stratified_split(schedules, makespans, train_ratio=0.9, seed=97)

    np.save(os.path.join(out_dir, "schedules_train.npy"), schedules[train_idx])
    np.save(os.path.join(out_dir, "makespans_train.npy"),  makespans[train_idx])
    np.save(os.path.join(out_dir, "schedules_val.npy"),   schedules[val_idx])
    np.save(os.path.join(out_dir, "makespans_val.npy"),    makespans[val_idx])
    np.save(os.path.join(out_dir, "min_makespan.npy"),     makespans[:, -1].min())

    with open(os.path.join(raw_data_dir, "metadata.json")) as f:
        meta = json.load(f)
    meta["nb_train_samples"] = len(train_idx)
    meta["nb_val_samples"]   = len(val_idx)
    meta["data_source"]      = raw_data_dir
    if extra_meta:
        meta.update(extra_meta)
    with open(os.path.join(out_dir, "metadata.json"), "w") as f:
        json.dump(meta, f, indent=4)

    for fname in ["pfsp_instance.npy", "neh_makespan.npy",
                  "cds_makespan.npy", "palmer_makespan.npy"]:
        shutil.copy(os.path.join(raw_data_dir, fname), os.path.join(out_dir, fname))


def process_initial_dataset(raw_data_dir, out_dir, cfg, seed=97):
    """
    Build the iter-0 dataset according to cfg.

    Modes:
      "worst"            – keep the worst_pct% schedules (highest makespan)
      "uniform"          – uniform random sample of pct% of all schedules
      "percentile_range" – keep schedules whose makespan is between the low_pct and
                           high_pct percentiles (inclusive on both ends)

    Returns (schedules, makespans) for the selected subset.
    """
    if os.path.exists(os.path.join(out_dir, "metadata.json")):
        print(f"  [data] already done – skipping {out_dir}")
        s = np.concatenate([
            np.load(os.path.join(out_dir, "schedules_train.npy")),
            np.load(os.path.join(out_dir, "schedules_val.npy")),
        ])
        m = np.concatenate([
            np.load(os.path.join(out_dir, "makespans_train.npy")),
            np.load(os.path.join(out_dir, "makespans_val.npy")),
        ])
        return s, m

    schedules = np.load(os.path.join(raw_data_dir, "schedules.npy"))
    makespans  = np.load(os.path.join(raw_data_dir, "makespans.npy"))
    ms         = makespans[:, -1]
    n_total    = len(ms)
    mode       = cfg["mode"]

    if mode == "worst":
        pct     = cfg["pct"]
        n_keep  = math.floor(pct / 100 * n_total)
        keep_idx = np.argsort(ms)[-n_keep:]   # highest makespans
        extra   = {"init_mode": mode, "pct": pct, "n_kept": n_keep}

    elif mode == "uniform":
        pct     = cfg["pct"]
        n_keep  = math.floor(pct / 100 * n_total)
        np.random.seed(seed)
        keep_idx = np.random.choice(n_total, size=n_keep, replace=False)
        extra   = {"init_mode": mode, "pct": pct, "n_kept": n_keep}

    elif mode == "percentile_range":
        low_pct, high_pct = cfg["low_pct"], cfg["high_pct"]
        lo = np.percentile(ms, low_pct)
        hi = np.percentile(ms, high_pct)
        keep_idx = np.where((ms >= lo) & (ms <= hi))[0]
        extra   = {"init_mode": mode, "low_pct": low_pct, "high_pct": high_pct,
                   "n_kept": len(keep_idx)}

    else:
        raise ValueError(f"Unknown init mode: {mode!r}")

    schedules = schedules[keep_idx]
    makespans  = makespans[keep_idx]

    print(f"  [data] {_config_key(cfg)}: kept {len(schedules)} / {n_total} schedules"
          f"  (makespan range [{makespans[:,-1].min():.4f}, {makespans[:,-1].max():.4f}])")

    _save_split(out_dir, schedules, makespans, raw_data_dir, extra_meta=extra)
    return schedules, makespans


def _swap_augment(schedules, pfsp_instance, n_swaps, seen_schedules, rng):
    """
    For each (schedule, makespan) pair, generate up to n_swaps new schedules by
    randomly swapping two job positions.  Only unique schedules (not already in
    seen_schedules) are returned.

    Returns (new_schedules_list, new_makespans_list) — plain Python lists of arrays.
    """
    norm_factor = float(np.sum(pfsp_instance))
    aug_s, aug_m = [], []

    for s in schedules:
        n_jobs = len(s)
        generated = set()
        attempts  = 0
        max_attempts = n_swaps * 10   # avoid infinite loop when n_jobs is small

        while len(generated) < n_swaps and attempts < max_attempts:
            attempts += 1
            i, j = rng.choice(n_jobs, size=2, replace=False)
            swapped = s.copy()
            swapped[i], swapped[j] = swapped[j], swapped[i]
            t = tuple(swapped.tolist())
            if t in seen_schedules or t in generated:
                continue
            generated.add(t)
            cumul = evaluate_makespan(pfsp_instance, swapped.tolist())
            norm  = [c / norm_factor for c in cumul]
            aug_s.append(swapped)
            aug_m.append(norm)

    return aug_s, aug_m


def augment_and_resplit(pool_schedules, pool_makespans,
                        recover_dir, out_dir, seen_schedules, raw_data_dir,
                        n_swaps_per_schedule=0):
    """
    1. Load Phase-2 captures from recover_dir.
    2. Keep only schedules not already in seen_schedules.
    3. Optionally generate n_swaps_per_schedule swap variants per new schedule.
    4. Append everything unique to the pool, redo 90/10 split, save.

    Returns updated (pool_schedules, pool_makespans).
    """
    cap_s = np.load(os.path.join(recover_dir, "captured_schedules.npy"))   # (K, 9) int32
    cap_m = np.load(os.path.join(recover_dir, "captured_makespans.npy"))   # (K, 9) float32

    # ── deduplicate direct captures ───────────────────────────────────────────
    new_s, new_m = [], []
    for s, m in zip(cap_s, cap_m):
        t = tuple(s.tolist())
        if t not in seen_schedules:
            seen_schedules.add(t)
            new_s.append(s)
            new_m.append(m)

    n_new_direct = len(new_s)
    print(f"  [data] {n_new_direct} new unique captures "
          f"({len(cap_s) - n_new_direct} duplicates dropped)")

    # ── swap augmentation ─────────────────────────────────────────────────────
    n_new_swaps = 0
    if n_swaps_per_schedule > 0 and n_new_direct > 0:
        pfsp_instance = np.load(os.path.join(raw_data_dir, "pfsp_instance.npy")).tolist()
        rng = np.random.default_rng(seed=42)
        aug_s, aug_m = _swap_augment(
            new_s, pfsp_instance,
            n_swaps=n_swaps_per_schedule,
            seen_schedules=seen_schedules,
            rng=rng,
        )
        for s, m in zip(aug_s, aug_m):
            seen_schedules.add(tuple(s.tolist()))
        n_new_swaps = len(aug_s)
        new_s.extend(aug_s)
        new_m.extend([np.array(m, dtype=np.float32) for m in aug_m])
        print(f"  [data] +{n_new_swaps} swap-augmented schedules")

    n_added = len(new_s)
    print(f"  [data] total added to pool: {n_added} "
          f"(pool size: {len(pool_schedules)} → {len(pool_schedules) + n_added})")
    if n_added > 0:
        pool_schedules = np.concatenate(
            [pool_schedules, np.stack(new_s).astype(np.int32)])
        pool_makespans  = np.concatenate(
            [pool_makespans,  np.stack(new_m).astype(np.float32)])

    if os.path.exists(os.path.join(out_dir, "metadata.json")):
        print(f"  [data] already done – skipping {out_dir}")
        return pool_schedules, pool_makespans

    _save_split(out_dir, pool_schedules, pool_makespans, raw_data_dir,
                extra_meta={"n_pool": len(pool_schedules),
                            "n_new_captures": n_new_direct,
                            "n_new_swaps": n_new_swaps})
    return pool_schedules, pool_makespans


def _count_new_captures(r_dir, seen_schedules):
    """Count how many schedules in r_dir are not already in seen_schedules (non-mutating)."""
    cap_s = np.load(os.path.join(r_dir, "captured_schedules.npy"))
    return sum(1 for s in cap_s if tuple(s.tolist()) not in seen_schedules)


# ── main loop ─────────────────────────────────────────────────────────────────

for cfg in (outer := tqdm(INIT_CONFIGS, desc="init_config")):
    cfg_key = _config_key(cfg)
    outer.set_description(cfg_key)
    base = os.path.join(XP3_OUT_ROOT, cfg_key)

    def data_dir(it):
        return os.path.join(base, f"data_iter_{it}")

    def model_dir(it):
        return os.path.join(base, f"model_iter_{it}")

    def recover_dir(it):
        return os.path.join(base, f"recover_iter_{it}")

    # iter 0: build initial dataset
    pool_s, pool_m = process_initial_dataset(RAW_DATA_DIR, data_dir(0), cfg)
    neh_makespan   = float(np.load(os.path.join(RAW_DATA_DIR, "neh_makespan.npy")))

    # seed seen_schedules with the initial pool so captures that duplicate it are dropped
    seen_schedules = set(map(tuple, pool_s.tolist()))

    for it in tqdm(range(N_ITERATIONS), desc="iterations", leave=False):
        best_so_far = pool_m[:, -1].min()
        print(f"\n[{cfg_key}] ── cycle {it + 1} / {N_ITERATIONS} "
              f"│ pool={len(pool_s)} schedules "
              f"│ best makespan={best_so_far:.6f} "
              f"│ NEH={neh_makespan:.6f} "
              f"│ gap={100 * (best_so_far - neh_makespan) / neh_makespan:+.2f}%"
              f" ──────────────")

        # ── Phase 1: train surrogate ──────────────────────────────────────────
        train(
            testing=False,
            seed=97,
            data_dir=data_dir(it),
            n_embd=N_EMBD,
            n_head=N_HEAD,
            n_layer=N_LAYER,
            intermediate_schedules=True,
            dropout=0.0,
            ff_width=4,
            train_batch_size=512,
            val_batch_size=256,
            nb_epochs=5,
            early_stopping_patience=15,
            checkpoint_interval_ratio=1.0,
            decay_lr=True,
            lr_partitions_ratios=[0.66],
            init_lr=1e-4,
            max_lr=1e-3,
            min_lr=5e-5,
            lr_warmup_iters_ratio=0.1,
            lr_decay_iters_ratio=0.95,
            beta1=0.9,
            beta2=0.95,
            weight_decay=10.0,
            grad_clip=1.0,
            compile=False,
            compile_mode="default",
            save_only_last_checkpoint=True,
            output_dir=model_dir(it),
        )

        # ── Phase 2: latent optimisation with schedule capture + retry ────────
        def _run_phase2(r_dir, seed):
            recover_schedules(
                testing=False,
                seed=seed,
                init_mode="random",
                n_optimization_steps=2000,
                epsilon=0.01,
                nb_sinkhorn_iters=40,
                decay_ls=False,
                ls_partitions_ratios=[0.66],
                init_ls=2.0,
                max_ls=10,
                min_ls=0.1,
                ls_warmup_iters_ratio=0.2,
                ls_decay_iters_ratio=0.90,
                decay_lr=True,
                lr_partitions_ratios=[0.66],
                init_lr=1e-2,
                max_lr=1e-1,
                min_lr=1e-4,
                lr_warmup_iters_ratio=0.1,
                lr_decay_iters_ratio=0.90,
                checkpoint_interval=1,
                data_dir=data_dir(it),
                model_dir=model_dir(it),
                output_dir=r_dir,
                apply_metropolized_filling=True,
                mf_patience=10,
                mf_fifo_size=50,
                mf_penalty_strength=0.25,
                mf_min_delta=0.0005,
                capture_schedules=True,
            )

        winning_recover_dir = None
        for attempt in range(MAX_RETRIES + 1):
            seed = 97 + attempt
            r_dir = (recover_dir(it) if attempt == 0
                     else f"{recover_dir(it)}_retry_{attempt}")
            _run_phase2(r_dir, seed)
            n_new = _count_new_captures(r_dir, seen_schedules)
            if n_new > 0:
                winning_recover_dir = r_dir
                break
            if attempt < MAX_RETRIES:
                print(f"  [retry] iter={it} attempt={attempt}: 0 new schedules, "
                      f"retrying with seed {seed + 1}")
        if winning_recover_dir is None:
            print(f"  [warn] iter={it}: no new schedules after {MAX_RETRIES + 1} attempts, "
                  f"proceeding without augmentation this cycle")

        # ── Augment pool + resplit for next iteration ─────────────────────────
        if it < N_ITERATIONS - 1 and winning_recover_dir is not None:
            prev_best = pool_m[:, -1].min()
            pool_s, pool_m = augment_and_resplit(
                pool_s, pool_m,
                winning_recover_dir, data_dir(it + 1),
                seen_schedules, RAW_DATA_DIR,
                n_swaps_per_schedule=N_SWAPS_PER_SCHEDULE,
            )
            new_best = pool_m[:, -1].min()
            improved = " ✓ improved" if new_best < prev_best else ""
            print(f"  [best] makespan after augmentation: {new_best:.6f}"
                  f" (was {prev_best:.6f}){improved}")