"""
Universal DRL model for CRMP: Train once, solve any instance instantly.

Train on thousands of random CRMP instances.
At inference: 5ms per new instance (vs GA's 1-2 seconds).
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import time
from itertools import permutations
from crmp_env import (CRMPEnv, evaluate_sequence, simulate_crmp,
                      NUM_JOBS_A, NUM_JOBS_B, NUM_MACHINES_A, NUM_MACHINES_B,
                      LINE_A_PROC, LINE_B_PROC,
                      LINE_A_YIELD_GRAN, LINE_A_YIELD_STRIP,
                      LINE_B_DEMAND_GRAN, LINE_B_DEMAND_STRIP)


class UniversalAgent(nn.Module):
    """Larger model for generalization across instances."""
    def __init__(self, obs_dim, hidden=256, latent=128):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, latent), nn.ReLU(),
        )
        self.policy_a = nn.Sequential(
            nn.Linear(latent, 128), nn.ReLU(),
            nn.Linear(128, NUM_JOBS_A + 1),
        )
        self.policy_b = nn.Sequential(
            nn.Linear(latent, 128), nn.ReLU(),
            nn.Linear(128, NUM_JOBS_B + 1),
        )
        self.value_head = nn.Sequential(
            nn.Linear(latent, 128), nn.ReLU(),
            nn.Linear(128, 1),
        )

    def forward(self, obs, mask_a=None, mask_b=None):
        z = self.encoder(obs)
        la = self.policy_a(z)
        lb = self.policy_b(z)
        if mask_a is not None:
            la = la + (1 - mask_a) * (-1e8)
        if mask_b is not None:
            lb = lb + (1 - mask_b) * (-1e8)
        return la, lb, self.value_head(z)


def generate_instance(rng, scale=(0.6, 1.4)):
    """Generate a random CRMP instance."""
    lo, hi = scale
    pa = np.maximum(LINE_A_PROC * rng.uniform(lo, hi, LINE_A_PROC.shape), 1.0)
    pb = np.maximum(LINE_B_PROC * rng.uniform(lo, hi, LINE_B_PROC.shape), 1.0)
    yg = np.maximum(LINE_A_YIELD_GRAN * rng.uniform(lo, hi, LINE_A_YIELD_GRAN.shape), 1.0)
    ys = np.maximum(LINE_A_YIELD_STRIP * rng.uniform(lo, hi, LINE_A_YIELD_STRIP.shape), 1.0)
    dg = LINE_B_DEMAND_GRAN * rng.uniform(lo, hi, LINE_B_DEMAND_GRAN.shape)
    ds = LINE_B_DEMAND_STRIP * rng.uniform(lo, hi, LINE_B_DEMAND_STRIP.shape)
    if dg.sum() > yg.sum() * 0.95:
        dg *= (yg.sum() * 0.95) / dg.sum()
    if ds.sum() > ys.sum() * 0.95:
        ds *= (ys.sum() * 0.95) / ds.sum()
    return pa, pb, yg, ys, dg, ds


def collect_episode(env, agent, device, deterministic=False):
    obs = env.reset()
    data = {'obs': [], 'mask_a': [], 'mask_b': [],
            'act_a': [], 'act_b': [],
            'logp_a': [], 'logp_b': [],
            'values': [], 'rewards': [], 'dones': []}
    done = False
    while not done:
        obs_t = torch.FloatTensor(obs).unsqueeze(0).to(device)
        ma = torch.FloatTensor(env.get_mask_a()).unsqueeze(0).to(device)
        mb = torch.FloatTensor(env.get_mask_b()).unsqueeze(0).to(device)
        with torch.no_grad():
            la, lb, val = agent(obs_t, ma, mb)
            da = torch.distributions.Categorical(logits=la)
            db = torch.distributions.Categorical(logits=lb)
            if deterministic:
                aa, ab = la.argmax(-1), lb.argmax(-1)
            else:
                aa, ab = da.sample(), db.sample()
        data['obs'].append(obs)
        data['mask_a'].append(ma.squeeze(0).cpu().numpy())
        data['mask_b'].append(mb.squeeze(0).cpu().numpy())
        data['act_a'].append(aa.item())
        data['act_b'].append(ab.item())
        data['logp_a'].append(da.log_prob(aa).item())
        data['logp_b'].append(db.log_prob(ab).item())
        data['values'].append(val.item())
        obs, reward, done, info = env.step(aa.item(), ab.item())
        data['rewards'].append(reward)
        data['dones'].append(done)
    return data, info


def compute_gae(rewards, values, dones, gamma=0.99, lam=0.95):
    advantages, gae, nv = [], 0, 0
    for t in reversed(range(len(rewards))):
        if dones[t]: nv, gae = 0, 0
        delta = rewards[t] + gamma * nv - values[t]
        gae = delta + gamma * lam * gae
        advantages.insert(0, gae)
        nv = values[t]
    returns = [a + v for a, v in zip(advantages, values)]
    return returns, advantages


def sa_solve(pa, pb, yg, ys, dg, ds, n_starts=10, max_iter=20000, seed=42):
    """SA baseline for comparison."""
    rng = np.random.default_rng(seed)
    all_b = list(permutations(range(NUM_JOBS_B)))
    results = []
    t0 = time.time()
    for s in range(n_starts):
        ca = rng.permutation(NUM_JOBS_A).tolist()
        cb = rng.permutation(NUM_JOBS_B).tolist()
        cms = simulate_crmp(ca, cb, pa, pb, yg, ys, dg, ds)["makespan"]
        ba, bb, bms = list(ca), list(cb), cms
        T = 80.0
        for i in range(max_iter):
            r = rng.random()
            na, nb = list(ca), list(cb)
            if r < 0.4:
                idx = rng.integers(len(na))
                v = na.pop(idx); na.insert(rng.integers(len(na)+1), v)
            elif r < 0.7:
                i1, i2 = rng.choice(len(na), 2, replace=False)
                na[i1], na[i2] = na[i2], na[i1]
            else:
                i1, i2 = rng.choice(len(nb), 2, replace=False)
                nb[i1], nb[i2] = nb[i2], nb[i1]
            nms = simulate_crmp(na, nb, pa, pb, yg, ys, dg, ds)["makespan"]
            d = nms - cms
            if d < 0 or rng.random() < np.exp(-d / max(T, 1e-10)):
                ca, cb, cms = na, nb, nms
                if cms < bms: ba, bb, bms = list(ca), list(cb), cms
            T *= 0.9997
        for perm in all_b:
            ms = simulate_crmp(ba, list(perm), pa, pb, yg, ys, dg, ds)["makespan"]
            if ms < bms: bms = ms
        results.append(bms)
    return {"best": min(results), "avg": np.mean(results),
            "std": np.std(results), "cpu": time.time() - t0}


def train():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    if device.type == 'cuda':
        print(f"GPU: {torch.cuda.get_device_name(0)}")

    # Get obs_dim from a dummy env
    dummy = CRMPEnv(stochastic=False)
    obs = dummy.reset()
    obs_dim = len(obs)

    agent = UniversalAgent(obs_dim).to(device)
    optimizer = torch.optim.Adam(agent.parameters(), lr=3e-4)

    num_epochs = 300
    eps_per_epoch = 128
    ent_coeff = 0.1
    rng = np.random.default_rng(42)

    best_real = float('inf')

    print(f"\n{'='*70}")
    print(f"Universal DRL Training for CRMP")
    print(f"Train on random instances, test on real + synthetic")
    print(f"Epochs: {num_epochs}, Episodes/epoch: {eps_per_epoch}")
    print(f"{'='*70}\n")

    t0 = time.time()

    for epoch in range(num_epochs):
        batch_obs, batch_ma, batch_mb = [], [], []
        batch_aa, batch_ab = [], []
        batch_lpa, batch_lpb = [], []
        batch_ret, batch_adv = [], []
        epoch_ms = []

        for _ in range(eps_per_epoch):
            # 80% random instances, 20% real instance
            if rng.random() < 0.8:
                pa, pb, yg, ys, dg, ds = generate_instance(rng)
            else:
                pa, pb = LINE_A_PROC, LINE_B_PROC
                yg, ys = LINE_A_YIELD_GRAN, LINE_A_YIELD_STRIP
                dg, ds = LINE_B_DEMAND_GRAN, LINE_B_DEMAND_STRIP

            env = CRMPEnv(stochastic=True, noise_std=0.02,
                          base_proc_a=pa, base_proc_b=pb,
                          base_yield_g=yg, base_yield_s=ys,
                          base_demand_g=dg, base_demand_s=ds)
            data, info = collect_episode(env, agent, device)
            ms = info.get('makespan') or 9999
            epoch_ms.append(ms)

            rets, advs = compute_gae(data['rewards'], data['values'], data['dones'])
            batch_obs.extend(data['obs'])
            batch_ma.extend(data['mask_a'])
            batch_mb.extend(data['mask_b'])
            batch_aa.extend(data['act_a'])
            batch_ab.extend(data['act_b'])
            batch_lpa.extend(data['logp_a'])
            batch_lpb.extend(data['logp_b'])
            batch_ret.extend(rets)
            batch_adv.extend(advs)

        # PPO update
        obs_t = torch.FloatTensor(np.array(batch_obs)).to(device)
        ma_t = torch.FloatTensor(np.array(batch_ma)).to(device)
        mb_t = torch.FloatTensor(np.array(batch_mb)).to(device)
        aa_t = torch.LongTensor(batch_aa).to(device)
        ab_t = torch.LongTensor(batch_ab).to(device)
        old_lpa = torch.FloatTensor(batch_lpa).to(device)
        old_lpb = torch.FloatTensor(batch_lpb).to(device)
        ret_t = torch.FloatTensor(batch_ret).to(device)
        adv_t = torch.FloatTensor(batch_adv).to(device)
        adv_t = (adv_t - adv_t.mean()) / (adv_t.std() + 1e-8)

        n = len(batch_obs)
        bs = min(512, n)
        idx_all = np.arange(n)
        for _ in range(6):
            np.random.shuffle(idx_all)
            for start in range(0, n, bs):
                idx = idx_all[start:min(start+bs, n)]
                la, lb, vals = agent(obs_t[idx], ma_t[idx], mb_t[idx])
                da = torch.distributions.Categorical(logits=la)
                db = torch.distributions.Categorical(logits=lb)
                nlpa = da.log_prob(aa_t[idx])
                nlpb = db.log_prob(ab_t[idx])
                ratio = torch.exp((nlpa - old_lpa[idx]) + (nlpb - old_lpb[idx]))
                s1 = ratio * adv_t[idx]
                s2 = torch.clamp(ratio, 0.8, 1.2) * adv_t[idx]
                ploss = -torch.min(s1, s2).mean()
                vloss = F.mse_loss(vals.squeeze(), ret_t[idx])
                ent = (da.entropy() + db.entropy()).mean()
                loss = ploss + 0.5*vloss - ent_coeff*ent
                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(agent.parameters(), 0.5)
                optimizer.step()

        # LR schedule
        lr = 3e-4 * max(0.05, 1 - epoch / num_epochs)
        for pg in optimizer.param_groups: pg['lr'] = lr
        if epoch > 100:
            ent_coeff = max(0.01, ent_coeff * 0.997)

        # Evaluate on real instance
        if (epoch + 1) % 10 == 0 or epoch < 10:
            real_env = CRMPEnv(stochastic=False)
            _, ri = collect_episode(real_env, agent, device, deterministic=True)
            real_ms = ri.get('makespan') or 9999

            # Sample 100 from real
            sample_best = real_ms
            for _ in range(100):
                se = CRMPEnv(stochastic=False)
                _, si = collect_episode(se, agent, device, deterministic=False)
                sms = si.get('makespan') or 9999
                if sms < sample_best: sample_best = sms

            if sample_best < best_real:
                best_real = sample_best
                torch.save(agent.state_dict(), 'universal_agent.pt')

            elapsed = time.time() - t0
            avg_ms = np.mean(epoch_ms)
            marker = " <<<MATCH/BEAT GA>>>" if sample_best <= 1307 else ""
            print(f"E{epoch+1:4d} | Real: det={real_ms:.0f} samp={sample_best:.0f} "
                  f"best={best_real:.0f} | Avg:{avg_ms:.0f} | {elapsed:.0f}s{marker}")

    train_time = time.time() - t0

    # ==================== EVALUATION ====================
    print(f"\n{'='*70}")
    print(f"EVALUATION (train time: {train_time:.0f}s)")
    print(f"{'='*70}")

    # Load best model
    agent.load_state_dict(torch.load('universal_agent.pt', weights_only=True))
    agent.eval()

    # --- Real dataset (Table 5) ---
    print("\n--- Table 5: Real Dataset ---")

    # DRL: deterministic + sampling
    real_env = CRMPEnv(stochastic=False)
    _, ri = collect_episode(real_env, agent, device, deterministic=True)
    drl_det = ri.get('makespan') or 9999

    drl_samples = []
    for _ in range(1000):
        se = CRMPEnv(stochastic=False)
        _, si = collect_episode(se, agent, device, deterministic=False)
        drl_samples.append(si.get('makespan') or 9999)

    # Inference speed
    t1 = time.time()
    for _ in range(1000):
        ie = CRMPEnv(stochastic=False)
        _, _ = collect_episode(ie, agent, device, deterministic=True)
    infer_ms = (time.time() - t1) / 1000 * 1000

    print(f"DRL deterministic:  {drl_det:.0f}")
    print(f"DRL best (1k samp): {min(drl_samples):.0f}")
    print(f"DRL avg (1k samp):  {np.mean(drl_samples):.1f}")
    print(f"DRL std:            {np.std(drl_samples):.2f}")
    print(f"DRL inference:      {infer_ms:.2f} ms/episode")

    # SA baseline
    print("\nRunning SA baseline on real data...")
    sa_real = sa_solve(LINE_A_PROC, LINE_B_PROC, LINE_A_YIELD_GRAN,
                       LINE_A_YIELD_STRIP, LINE_B_DEMAND_GRAN,
                       LINE_B_DEMAND_STRIP, n_starts=10, max_iter=20000)
    print(f"SA best:  {sa_real['best']:.0f}")
    print(f"SA avg:   {sa_real['avg']:.1f}")
    print(f"SA std:   {sa_real['std']:.2f}")
    print(f"SA cpu:   {sa_real['cpu']:.2f}s")

    print(f"\n{'Method':<18} {'Best':>6} {'Avg':>8} {'Std':>8} {'Time':>12}")
    print("-" * 54)
    print(f"{'FCFS':<18} {'1438':>6} {'1438':>8} {'—':>8} {'—':>12}")
    print(f"{'Paper GA':<18} {'1307':>6} {'1315':>8} {'8.05':>8} {'1.28s':>12}")
    print(f"{'SA (ours)':<18} {sa_real['best']:>6.0f} {sa_real['avg']:>8.1f} {sa_real['std']:>8.2f} {sa_real['cpu']:>10.2f}s")
    print(f"{'DRL (ours)':<18} {min(drl_samples):>6.0f} {np.mean(drl_samples):>8.1f} {np.std(drl_samples):>8.2f} {infer_ms:>8.2f}ms")
    print(f"{'Speedup':<18} {'':>6} {'':>8} {'':>8} {sa_real['cpu']/(infer_ms/1000):>8.0f}x")

    # --- Synthetic dataset (Table 6) ---
    print(f"\n--- Table 6: Synthetic Dataset (10 instances) ---")
    t6_sa, t6_drl, t6_fcfs = [], [], []
    sa_times, drl_times = [], []

    for inst in range(10):
        pa, pb, yg, ys, dg, ds = generate_instance(
            np.random.default_rng(inst*100+7))

        # FCFS
        f = simulate_crmp(list(range(8)), list(range(6)), pa, pb, yg, ys, dg, ds)["makespan"]
        t6_fcfs.append(f)

        # SA
        sa = sa_solve(pa, pb, yg, ys, dg, ds, n_starts=5, max_iter=15000, seed=inst)
        t6_sa.append(sa['best'])
        sa_times.append(sa['cpu'])

        # DRL (just inference - no retraining!)
        t_drl = time.time()
        drl_best = float('inf')
        drl_all = []
        for _ in range(300):
            ie = CRMPEnv(stochastic=False, base_proc_a=pa, base_proc_b=pb,
                         base_yield_g=yg, base_yield_s=ys,
                         base_demand_g=dg, base_demand_s=ds)
            _, si = collect_episode(ie, agent, device, deterministic=False)
            ms = si.get('makespan') or 9999
            drl_all.append(ms)
            if ms < drl_best: drl_best = ms
        drl_cpu = time.time() - t_drl
        t6_drl.append(drl_best)
        drl_times.append(drl_cpu)

        print(f"  Inst {inst+1:2d}: FCFS={f:.0f} SA={sa['best']:.0f}({sa['cpu']:.1f}s) "
              f"DRL={drl_best:.0f}({drl_cpu:.1f}s)")

    print(f"\n{'Inst':<6} {'FCFS':>8} {'SA':>8} {'DRL':>8}")
    print("-" * 32)
    for i in range(10):
        best_mark = " *" if t6_drl[i] <= t6_sa[i] else ""
        print(f"{'#'+str(i+1):<6} {t6_fcfs[i]:>8.0f} {t6_sa[i]:>8.0f} {t6_drl[i]:>8.0f}{best_mark}")
    print("-" * 32)
    print(f"{'Avg':<6} {np.mean(t6_fcfs):>8.0f} {np.mean(t6_sa):>8.0f} {np.mean(t6_drl):>8.0f}")

    wins = sum(1 for d, s in zip(t6_drl, t6_sa) if d <= s)
    print(f"\nDRL wins/ties: {wins}/10")
    print(f"SA avg time:  {np.mean(sa_times):.1f}s per instance")
    print(f"DRL avg time: {np.mean(drl_times):.1f}s (300 samples)")
    print(f"DRL 1-shot:   {infer_ms:.2f}ms")

    print(f"\n{'='*70}")
    print(f"SUMMARY")
    print(f"  Training: {train_time:.0f}s (one-time cost)")
    print(f"  Real data: DRL best={min(drl_samples):.0f} vs GA=1307")
    print(f"  Synthetic: DRL avg={np.mean(t6_drl):.0f} vs SA avg={np.mean(t6_sa):.0f}")
    print(f"  Speed: {infer_ms:.2f}ms vs SA {np.mean(sa_times):.1f}s = {np.mean(sa_times)/(infer_ms/1000):.0f}x faster")
    print(f"{'='*70}")


if __name__ == '__main__':
    train()