"""Fresh-trajectory clipped PPO and weak offline supervision."""
import random

import torch
from torch.nn import functional as F

from qojbot import rules as R
from qojbot.history import latest_log
from qojbot.model_benchmark import play_game

from .features import collate
from .network import score_loss
from .storage import cpu_tree


def terminal_targets(samples, device):
    rewards = torch.tensor([sample['reward'] for sample in samples], device=device)
    return rewards, rewards / 6.


def losses(model, samples, device, ppo=False, anchor_weight=.05,
           history_weight=None, distill_weight=0.):
    batch = collate(samples, device)
    output = model(batch)
    # Batch the candidate softmax to avoid tiny MPS dispatches per state.
    counts = [len(sample['actions']) for sample in samples]
    cursor = sum(counts)
    rows = torch.arange(len(samples), device=device)
    local = torch.cat([torch.arange(count) for count in counts]).to(device)
    base = torch.cat([sample['base'] for sample in samples]).to(device)
    logits = output['prior_scale'][batch['owner']] * base + output['correction']
    shape = (len(samples), max(counts))
    padded = torch.full(shape, -torch.inf, device=device)
    padded[batch['owner'], local] = logits
    teacher_padded = torch.full(shape, -torch.inf, device=device)
    teacher_padded[batch['owner'], local] = base
    valid = torch.arange(max(counts), device=device)[None] < torch.tensor(counts, device=device)[:, None]
    log_probs = padded.log_softmax(-1)
    probs = log_probs.exp()
    teacher_logp = teacher_padded.log_softmax(-1)
    teacher = teacher_logp.exp()
    safe_logp = log_probs.masked_fill(~valid, 0)
    safe_teacher_logp = teacher_logp.masked_fill(~valid, 0)
    anchor = (teacher * (safe_teacher_logp - safe_logp)).sum(-1).mean()
    entropy = -(probs * safe_logp).sum(-1).mean()
    chosen = torch.tensor([sample['selected'] for sample in samples], device=device)
    new_logp = log_probs[rows, chosen]
    approx_kl = new_logp.sum() * 0
    if ppo:
        old_logp = torch.tensor([sample['old_logp'] for sample in samples], device=device)
        advantage = torch.tensor([sample['advantage'] for sample in samples], device=device)
        ratio = (new_logp - old_logp).exp()
        policy_loss = -torch.minimum(ratio * advantage, ratio.clamp(.8, 1.2) * advantage).mean()
        approx_kl = ((ratio - 1) - (new_logp - old_logp)).mean()
    else:
        weights = torch.tensor([history_weight if history_weight is not None and
                                sample.get('origin') == 'history' else sample.get('weight', .15)
                                for sample in samples], device=device)
        policy_loss = -(new_logp * weights).mean()
        if distill_weight:
            # Reward the frozen strong teacher's preferred action, never the
            # unfiltered human action. KL keeps the rest of its ranking intact.
            preferred = teacher_padded.argmax(-1)
            policy_loss = policy_loss - distill_weight * log_probs[rows, preferred].mean()
    offsets, offset = [], 0
    for count, sample in zip(counts, samples):
        offsets.append(offset + sample['selected'])
        offset += count
    selected = torch.tensor(offsets, device=device)
    rewards, transformed = terminal_targets(samples, device)
    # Historical returns describe another behavior policy; use histories for
    # public-state distillation/beliefs, actual on-policy/teacher games for value.
    trusted = [i for i, s in enumerate(samples) if history_weight != 0 or s.get('origin') != 'history']
    zero = output['value'].sum() * 0 + output['score_logits'].sum() * 0 + output['win_logits'].sum() * 0
    value_loss = action_value_loss = win_loss = zero
    if trusted:
        value_loss = F.smooth_l1_loss(output['value'][trusted], transformed[trusted])
        action_value_loss = score_loss(output['score_logits'][selected[trusted]], rewards[trusted], model.support)
        win_loss = F.binary_cross_entropy_with_logits(output['win_logits'][selected[trusted]],
                                                     (rewards[trusted] > 0).float())
    belief_indices = [i for i, sample in enumerate(samples) if 'belief_target' in sample and sample['role'] != 3]
    belief_loss = output['value'].sum() * 0
    if belief_indices:
        targets = torch.stack([samples[i]['belief_target'] for i in belief_indices]).to(device)
        # Counts, not rounded hand assignments: supervision is inaccessible at inference.
        belief_loss = F.smooth_l1_loss(output['belief_counts'][belief_indices], targets)
    loss = (policy_loss + anchor_weight * anchor + .5 * value_loss +
            .15 * action_value_loss + .05 * win_loss + .05 * belief_loss -
            (.01 * entropy if ppo else 0))
    stats = {'loss': float(loss.detach().cpu()), 'policy': float(policy_loss.detach().cpu()),
             'anchor_kl': float(anchor.detach().cpu()), 'value': float(value_loss.detach().cpu()),
             'belief': float(belief_loss.detach().cpu()),
             'approx_kl': float(approx_kl.detach().cpu()) if ppo else 0.,
             'samples': len(samples), 'actions': cursor}
    return loss, stats


def update(model, optimizer, samples, device, ppo=False, anchor_weight=.05,
           history_weight=None, distill_weight=0.):
    model.train()
    loss, stats = losses(model, samples, device, ppo, anchor_weight, history_weight, distill_weight)
    if not torch.isfinite(loss):
        raise RuntimeError('nonfinite training loss')
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 5.)
    if not torch.isfinite(norm):
        raise RuntimeError('nonfinite training gradients')
    optimizer.step()
    return stats


def batches(samples, batch_states, max_actions=4096):
    pending, actions = [], 0
    for sample in samples:
        count = len(sample['actions'])
        if pending and (len(pending) >= batch_states or actions + count > max_actions):
            yield pending
            pending, actions = [], 0
        pending.append(sample)
        actions += count
    if pending:
        yield pending


class Actor:
    def __init__(self, policy, rng, teacher=False):
        self.policy, self.rng, self.teacher = policy, rng, teacher
        self.pending = []

    def select(self, candidates, sample, state):
        if self.teacher:
            logits, value = sample['base'], 0.
            selected = int(logits.argmax())
        else:
            logits, value = self.policy.predict(sample)
            probs = logits.softmax(0).tolist()
            selected = self.rng.choices(range(len(probs)), weights=probs)[0]
        sample.update(selected=selected, old_logp=float(logits.log_softmax(0)[selected]),
                      old_value=value, weight=1., origin='synthetic')
        if not self.teacher:
            sample['old_log_probs'] = logits.log_softmax(0)
        attempt = sum(event.get('kind') == 'redeal' for event in state.get('log', []))
        self.pending.append((sample, state['seat'], attempt))
        return candidates[selected]

    def choose(self, obs, state):
        candidates, sample = self.policy.prepare(obs, state)
        return self.select(candidates, sample, state)

    def choose_bid(self, state):
        candidates, sample = self.policy.prepare_bid(state)
        return self.select(candidates, sample, state)

    def finish(self, game):
        hands = [list(hand) for hand in game['initial_hands']]
        snapshots = {}
        step = 0
        for event in latest_log(game):
            if event.get('kind') in ('play', 'pass'):
                snapshots[step] = [list(hand) for hand in hands]
                step += 1
                if event['kind'] == 'play':
                    for card in event['cards']:
                        hands[event['seat']].remove(card)
        from .dataset import belief_labels
        samples, play_index = [], 0
        final_attempt = game['redeals']
        for sample, seat, attempt in self.pending:
            if attempt != final_attempt:
                continue
            sample['reward'] = float(game['deltas'][seat])
            # Actor can occupy multiple seats. Map samples to physical play events.
            if sample['role'] != 3:
                events = [event for event in latest_log(game)
                          if event.get('kind') in ('play', 'pass')]
                while play_index < len(events) and events[play_index]['seat'] != seat:
                    play_index += 1
                if play_index < len(events):
                    belief_labels(sample, snapshots[play_index], seat, game['landlord'])
                    play_index += 1
            samples.append(sample)
        self.pending.clear()
        return samples


class SearchOpponent:
    def __init__(self, seconds=.03):
        from qojbot.policy import SearchPolicy
        self.policy = SearchPolicy(seconds=seconds, samples=4)

    def choose(self, obs, state):
        return self.policy.choose(obs)

    def choose_bid(self, state):
        from qojbot.policy import assess_bid
        return assess_bid(R.from_cards(state['hand']), int(state.get('bid') or 0),
                          bool(state.get('must_bid')))['value']


def collect(policy, games, index, seed, rng, mode='bidding', league=(), opponent=None,
            targeted=False):
    """Current learner vs teacher, snapshot league, search and self-play."""
    all_samples, records = [], []
    target = opponent
    for offset in range(games):
        number = index + offset
        learner = Actor(policy, rng)
        kind = number % (8 if targeted else 6)
        landlord = (number // 24) % 3 if targeted else number % 3
        relative = (number // 8) % 3 if targeted else (number // 3) % 3
        singleton = (landlord + relative) % 3
        if kind == 0 and not targeted:
            opponent = SearchOpponent()
        elif ((kind == 7 if targeted else kind == 2) and league):
            from .policy import QOJPolicy
            opponent = QOJPolicy(device=str(policy.device), teacher=policy.teacher,
                                 config=policy.model.config)
            opponent.model.load_state_dict(league[rng.randrange(len(league))])
        else:
            opponent = target or policy.teacher
        lineup = ['learner'] * 3
        if kind < (3 if targeted else 2):
            lineup = ['opponent'] * 3
            lineup[singleton] = 'learner'
        elif (kind in (3, 4, 5, 7) if targeted else kind < 4):
            lineup[singleton] = 'opponent'
        record = play_game({'learner': learner, 'opponent': opponent}, lineup, seed + number,
                           first_bidder=landlord, mode=mode)
        all_samples.extend(learner.finish(record))
        records.append(record)
        print(f'collected game {number + 1}, turns={record["turns"]}, opponent={type(opponent).__name__}', flush=True)
    advantages = torch.tensor([s['reward'] / 6. - s['old_value']
                              for s in all_samples])
    # Role-wise scaling avoids the landlord's x2 score dominating farmer updates.
    for role in range(4):
        indices = [i for i, s in enumerate(all_samples) if s['role'] == role]
        if indices:
            local = advantages[indices]
            std = local.std(unbiased=False).clamp_min(.5)
            local = (local - local.mean()) / std if len(indices) > 1 else local / std
            for i, advantage in zip(indices, local.tolist()):
                all_samples[i]['advantage'] = advantage
    return all_samples, records


def rollout_kl(policy, samples, batch_states=32, max_actions=4096):
    """Exact old||new KL over every legal action on the complete rollout."""
    total, count = 0., 0
    policy.model.eval()
    with torch.no_grad():
        for samples_batch in batches(samples, batch_states, max_actions):
            batch = collate(samples_batch, policy.device)
            output = policy.model(batch)
            counts = [len(s['actions']) for s in samples_batch]
            local = torch.cat([torch.arange(n) for n in counts]).to(policy.device)
            base = torch.cat([s['base'] for s in samples_batch]).to(policy.device)
            logits = output['prior_scale'][batch['owner']] * base + output['correction']
            shape = (len(counts), max(counts))
            padded = torch.full(shape, -torch.inf, device=policy.device)
            padded[batch['owner'], local] = logits
            old = torch.full(shape, -torch.inf, device=policy.device)
            old[batch['owner'], local] = torch.cat([s['old_log_probs'] for s in samples_batch]).to(policy.device)
            valid = torch.arange(max(counts), device=policy.device)[None] < torch.tensor(counts, device=policy.device)[:, None]
            difference = old.masked_fill(~valid, 0) - padded.log_softmax(-1).masked_fill(~valid, 0)
            total += float((old.exp() * difference).sum().cpu())
            count += len(counts)
    return max(0., total / max(count, 1))


def warmup(policy, games, seed, rng):
    actor = Actor(policy, rng, teacher=True)
    samples, records = [], []
    for index in range(games):
        record = play_game({'teacher': actor}, ['teacher'] * 3, seed + index,
                           first_bidder=index % 3, mode='bidding')
        samples.extend(actor.finish(record))
        records.append(record)
        print(f'teacher warmup {index + 1}/{games}', flush=True)
    return samples, records