"""
Reward-Shaped Cost Transport (RSCT) — RL for PRZELOM.

Uses the attention EOT cost matrix as a reward-shaped optimization surface.
Iteratively refines the transport plan using per-example evaluation rewards.
C_shaping -= lr * advantage * T_current (REINFORCE on transport plan).
"""
from __future__ import annotations
import json
import math
import os
import tempfile
from collections import defaultdict
from datetime import datetime
from typing import Dict, List, Tuple
import torch
from wisent.core.control.steering_methods.configs.optimal import get_optimal
from wisent.core.control.steering_methods.steering_object import SteeringObjectMetadata
from wisent.core.control.steering_methods.methods.przelom.przelom_steering_object import (
    PrzelomSteeringObject,
)
from wisent.core.control.steering_methods.methods.przelom.create import _regularized_pinv
from wisent.core.control.steering_methods.methods.szlak.transport import (
    compute_attention_affinity_cost,
)
from wisent.core.utils.cli.optimize_steering.data.responses import execute_generate_responses
from wisent.core.utils.cli.optimize_steering.scores import execute_evaluate_responses
from wisent.core.utils.cli.optimize_steering.pipeline import _make_args
from wisent.core.utils import preferred_dtype
from wisent.core.utils.config_tools.constants import (LOG_EPS,
    SEPARATOR_WIDTH_WIDE)


def _precompute_layer_data(
    data: dict, epsilon: float,
) -> Tuple[Dict[int, dict], SteeringObjectMetadata]:
    """One-time C_original / q_neg / k_pos / T_current computation per layer."""
    dtype = preferred_dtype()
    pairs_list = data.get('pairs', [])
    num_heads = data.get('num_attention_heads')
    num_kv_heads = data.get('num_key_value_heads')
    layer_act = defaultdict(lambda: {
        "positive": [], "negative": [],
        "q_proj_activations": [], "k_proj_activations": [],
    })
    for pair in pairs_list:
        for ls, al in pair['positive_response'].get('layers_activations', {}).items():
            if al is not None:
                layer_act[ls]["positive"].append(torch.tensor(al, dtype=dtype))
        for ls, al in pair['negative_response'].get('layers_activations', {}).items():
            if al is not None:
                layer_act[ls]["negative"].append(torch.tensor(al, dtype=dtype))
        for ls, qv in pair['negative_response'].get('q_proj_activations', {}).items():
            if qv is not None:
                layer_act[ls]["q_proj_activations"].append(torch.tensor(qv, dtype=dtype))
        for ls, kv in pair['positive_response'].get('k_proj_activations', {}).items():
            if kv is not None:
                layer_act[ls]["k_proj_activations"].append(torch.tensor(kv, dtype=dtype))

    all_layers = sorted(layer_act.keys(), key=lambda x: int(x))
    hidden_dim = layer_act[all_layers[0]]["positive"][0].shape[-1]
    trait_label = data.get('trait_label', 'unknown')
    parts = trait_label.split('/')
    cal_raw = data.get('calibration_norms', {})
    metadata = SteeringObjectMetadata(
        method="przelom", model_name=data.get('model', 'unknown'),
        benchmark=parts[-1], category=parts[0] if len(parts) > 1 else 'unknown',
        extraction_strategy=data.get('token_aggregation', 'unknown'),
        num_pairs=len(pairs_list), layers=[int(l) for l in all_layers],
        hidden_dim=hidden_dim, created_at=datetime.now().isoformat(),
        calibration_norms={int(k): float(v) for k, v in cal_raw.items()},
        extraction_component=data.get('extraction_component', get_optimal("extraction_component")),
        extra={'num_attention_heads': num_heads, 'num_key_value_heads': num_kv_heads},
    )
    layer_data: Dict[int, dict] = {}
    for layer_str in all_layers:
        la = layer_act[layer_str]
        if not la["positive"] or not la["negative"]:
            continue
        if not la["q_proj_activations"] or not la["k_proj_activations"]:
            print(f"   Skipping layer {layer_str}: no Q/K projections")
            continue
        pos = torch.stack([t.detach().float().reshape(-1) for t in la["positive"]], dim=0)
        neg = torch.stack([t.detach().float().reshape(-1) for t in la["negative"]], dim=0)
        q_neg = torch.stack([t.detach().float().reshape(-1) for t in la["q_proj_activations"]], dim=0)
        k_pos = torch.stack([t.detach().float().reshape(-1) for t in la["k_proj_activations"]], dim=0)
        C = compute_attention_affinity_cost(q_neg, k_pos, num_heads=num_heads, num_kv_heads=num_kv_heads)
        T_current = torch.softmax(-C / epsilon, dim=1)
        li = int(layer_str)
        layer_data[li] = {
            "C_original": C, "T_current": T_current,
            "q_neg": q_neg, "k_pos": k_pos, "neg": neg, "pos": pos,
        }
        print(f"   Layer {layer_str}: n_neg={neg.shape[0]}, n_pos={pos.shape[0]}, "
              f"C_range=[{C.min().item():.3f}, {C.max().item():.3f}]")
    return layer_data, metadata


def _update_cost_shaping(
    C_shaping: Dict[int, torch.Tensor],
    layer_data: Dict[int, dict],
    rewards: List[float],
    learning_rate: float,
    baseline: float = None,
) -> None:
    """REINFORCE update: scalar advantage from mean reward shapes cost.

    The transport plan T_current (n_pairs x n_pairs) operates on enriched
    pairs, while rewards come from task evaluation (different examples).
    We use the aggregate reward as a scalar advantage: when steering performs
    above baseline, reinforce current transport paths; below, discourage them.
    """
    if not rewards:
        return
    if baseline is None:
        raise ValueError("baseline is required")
    mean_reward = sum(rewards) / len(rewards)
    advantage = mean_reward - baseline
    for layer_int, data in layer_data.items():
        T_current = data["T_current"]
        C_shaping[layer_int] -= learning_rate * advantage * T_current


def _create_przelom_from_shaping(
    metadata: SteeringObjectMetadata,
    layer_data: Dict[int, dict],
    C_shaping: Dict[int, torch.Tensor],
    epsilon: float,
    regularization: float,
    inference_k: int,
) -> PrzelomSteeringObject:
    """Create steering object from current shaped cost matrices."""
    num_heads = metadata.extra.get('num_attention_heads')
    num_kv_heads = metadata.extra.get('num_key_value_heads')
    source_points, disps = {}, {}
    for layer_int, data in layer_data.items():
        C_eff = data["C_original"] + C_shaping[layer_int]
        T_target = torch.softmax(-C_eff / epsilon, dim=1)
        T_current = data["T_current"]
        delta_C = epsilon * (torch.log(T_target.clamp(min=LOG_EPS))
                             - torch.log(T_current.clamp(min=LOG_EPS)))
        q_dim = data["q_neg"].shape[-1]
        k_dim = data["k_pos"].shape[-1]
        if q_dim == k_dim:
            k_pos_pinv = _regularized_pinv(data["k_pos"], regularization)
            delta_q = -math.sqrt(q_dim) * (delta_C @ k_pos_pinv.T)
        else:
            head_dim = q_dim // num_heads
            groups = num_heads // num_kv_heads
            k_by_head = data["k_pos"].reshape(-1, num_kv_heads, head_dim)
            delta_q_parts = []
            for g in range(num_kv_heads):
                k_g = k_by_head[:, g, :]
                k_g_pinv = _regularized_pinv(k_g, regularization)
                dq_g = -math.sqrt(head_dim) * (delta_C @ k_g_pinv.T)
                for _ in range(groups):
                    delta_q_parts.append(dq_g)
            delta_q = torch.cat(delta_q_parts, dim=-1)
        source_points[layer_int] = data["neg"].detach()
        disps[layer_int] = delta_q.detach()
    return PrzelomSteeringObject(
        metadata=metadata, source_points=source_points,
        displacements=disps, inference_k=inference_k,
    )


def _extract_per_example_rewards(scores_file: str) -> List[float]:
    """Extract per-example binary rewards from evaluation scores file."""
    with open(scores_file) as f:
        scores_data = json.load(f)
    rewards = []
    for ev in scores_data.get("evaluations", []):
        evaluation = ev.get("evaluation", {})
        if isinstance(evaluation, dict) and "correct" in evaluation:
            rewards.append(1.0 if evaluation["correct"] else 0.0)
        elif isinstance(evaluation, dict):
            rewards.append(evaluation["score"])
        else:
            rewards.append(0.0)
    return rewards


def _extract_aggregate_score(scores_file: str) -> float:
    """Extract aggregate accuracy from scores file."""
    with open(scores_file) as f:
        sd = json.load(f)
    return sd.get("aggregated_metrics", {}).get("acc") or sd.get("accuracy") or sd.get("acc") or 0.0


def execute_transport_rl(args):
    """Execute RL-based steering optimization. Dispatches by method.

    Transport methods (przelom, szlak): cost-matrix shaping with REINFORCE.
    All other methods: evolutionary strategy on steering vectors.
    """
    method = getattr(args, 'method', 'przelom').lower()
    if method not in ('przelom', 'szlak'):
        from wisent.core.utils.cli.optimize_steering.transport.vector_rl import run_vector_rl_loop
        return run_vector_rl_loop(args)

    # --- Transport cost-shaping path (PRZELOM / SZLAK) ---
    model = args.model
    task = args.task
    epf = args.enriched_pairs_file
    layer_sweep_strength = getattr(args, 'layer_sweep_strength', None)
    if layer_sweep_strength is None:
        raise ValueError("layer_sweep_strength is required")
    max_iter = getattr(args, 'max_iterations', None)
    if max_iter is None:
        raise ValueError("max_iterations is required")
    lr = getattr(args, 'learning_rate', None)
    if lr is None:
        raise ValueError("learning_rate is required")
    epsilon = args.epsilon
    reg = getattr(args, 'regularization', None)
    inf_k = getattr(args, 'inference_k', None)
    limit = getattr(args, 'limit', None)
    if limit is None:
        raise ValueError("limit is required")
    initial_baseline = getattr(args, 'initial_baseline', None)
    if initial_baseline is None:
        raise ValueError("initial_baseline is required")
    output_path = getattr(args, 'output', 'best_transport_rl.pt')
    device = getattr(args, 'device', None)

    print(f"\n{'=' * SEPARATOR_WIDTH_WIDE}")
    print(f"REWARD-SHAPED COST TRANSPORT (RSCT)")
    print(f"{'=' * SEPARATOR_WIDTH_WIDE}")
    print(f"   Model: {model}  |  Task: {task}")
    print(f"   Iterations: {max_iter}  |  LR: {lr}  |  Epsilon: {epsilon}")
    print(f"   Regularization: {reg}  |  Inference K: {inf_k}  |  Limit: {limit}")
    print(f"{'=' * SEPARATOR_WIDTH_WIDE}\n")

    # 1. Load and precompute
    print("Loading enriched pairs and precomputing layer data...")
    with open(epf, 'r') as f:
        data = json.load(f)
    layer_data, metadata = _precompute_layer_data(data, epsilon)
    if not layer_data:
        raise ValueError("No layers with Q/K projections found in enriched pairs")

    # 2. Initialize cost shaping
    C_shaping: Dict[int, torch.Tensor] = {
        li: torch.zeros_like(ld["C_original"]) for li, ld in layer_data.items()
    }
    best_score = -float('inf')
    best_obj = None
    score_history: List[float] = []

    # 3. RL loop
    for it in range(1, max_iter + 1):
        print(f"\n--- Iteration {it}/{max_iter} ---")
        with tempfile.TemporaryDirectory() as wd:
            sf = os.path.join(wd, "steering.pt")
            rf = os.path.join(wd, "responses.json")
            ef = os.path.join(wd, "scores.json")
            # Create steering object from shaped cost
            obj = _create_przelom_from_shaping(metadata, layer_data, C_shaping, epsilon, reg, inf_k)
            obj.save(sf)
            # Generate steered responses
            execute_generate_responses(_make_args(
                task=task, input_file=epf, model=model, output=rf,
                num_questions=limit, min_load_limit_questions=limit,
                steering_object=sf, steering_strength=layer_sweep_strength,
                steering_strategy=get_optimal("steering_strategy"), use_steering=True, device=device,
                verbose=False, cached_model=None,
            ))
            # Evaluate
            execute_evaluate_responses(_make_args(input=rf, output=ef, task=task, verbose=False))
            rewards = _extract_per_example_rewards(ef)
            agg = _extract_aggregate_score(ef)
            n = max(len(rewards), 1)
            print(f"   Score: {agg:.4f} (mean_reward={sum(rewards)/n:.3f}, n={len(rewards)})")
            # Update cost shaping with running baseline
            if rewards:
                baseline = sum(score_history) / len(score_history) if score_history else initial_baseline
                _update_cost_shaping(C_shaping, layer_data, rewards, lr, baseline=baseline)
            score_history.append(agg)
            # Track best — save incrementally so progress survives timeouts
            if agg > best_score:
                best_score = agg
                best_obj = obj
                os.makedirs(os.path.dirname(os.path.abspath(output_path)) or '.', exist_ok=True)
                best_obj.save(output_path)
                print(f"   ** New best: {best_score:.4f} (saved to {output_path})")

    # 4. Final save confirmation
    if best_obj is not None:
        best_obj.save(output_path)
        print(f"\nBest steering object saved to: {output_path} (score: {best_score:.4f})")
    else:
        print("\nNo valid steering object produced.")
    return {"best_score": best_score, "output": output_path}
