"""
Trust Proof v10: Multi-Scale Context Window + EMA Evaluation + L1 Loss

The canonical proof file. trust_v9's protocol verbatim (same data generation,
architecture, split, metric, EMA evaluation, checkpoints, plots) with ONE
change adopted after the Phase 26 attention-gap probe: criterion is
nn.L1Loss() instead of nn.MSELoss().

Why: the probe showed 94% of v9 HARD error was "attention right, pick wrong".
The cross-attention ledger identified the true expert on 99.4% of test tokens
but the output hedged 10-15% toward the runner-up source's value, because
MSE's optimum under a split belief is the mean. L1's optimum is the median,
which commits fully to the majority side of the belief.

Invariant preserved: the loss is a function of (model output, true value)
ONLY. No expert index, no source values, no boundary structure enters the
gradient. Same supervision as v9, different price curve.

L1 viability run (HARD, 60 epochs, seed 42, otherwise this protocol):
best 0.973 vs MSE 0.880 at matched epochs (MSE HARD benchmark: 0.900).
Attention-right errors' alpha median 0.865 -> 0.987, mid-blend 30.5% -> 18.4%,
ledger argmax == expert 0.991 (MSE: 0.994).

Imports only from time_transformer_proof.py (the timestamp encoding proof).
Everything else is self-contained.

Usage:
  python3 trust_v10.py              # HARD + SIMPLE, full models + ablations
  python3 trust_v10.py --hard       # HARD only, full model
  python3 trust_v10.py --simple     # SIMPLE only, full model
  python3 trust_v10.py --full-only  # both full models, no ablations
  python3 trust_v10.py --seed 123   # different seed
"""

import sys
sys.argv.append('--headless')

import math
import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, RandomSampler

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from time_transformer_proof import (
    NormalizedTimestampEncoding, TemporalAttention, TransformerLayer,
    inspect_encoding, device
)

PLOT_DIR = os.path.join(os.path.dirname(__file__), 'plots')
os.makedirs(PLOT_DIR, exist_ok=True)
CKPT_DIR = os.path.join(os.path.dirname(__file__), 'checkpoints')
os.makedirs(CKPT_DIR, exist_ok=True)


# ─── Config ───

SEED = 42
NUM_EPOCHS = 200
NUM_EPOCHS_ABLATION = 50
BATCH_SIZE = 32
K_TIMESTAMPS = 8
NUM_SAMPLES = 400000
BATCH_PER_EPOCH = 375
EMA_DECAY = 0.999

for i, arg in enumerate(sys.argv):
    if arg == '--seed' and i + 1 < len(sys.argv): SEED = int(sys.argv[i + 1])
    if arg == '--epochs' and i + 1 < len(sys.argv): NUM_EPOCHS = int(sys.argv[i + 1])
    if arg == '--samples' and i + 1 < len(sys.argv): NUM_SAMPLES = int(sys.argv[i + 1])


# ─── EMA ───

def clone_state(model_or_state):
    state = model_or_state.state_dict() if isinstance(model_or_state, nn.Module) else model_or_state
    return {k: v.detach().clone() for k, v in state.items()}


def ema_update(model, ema, decay):
    with torch.no_grad():
        for k, v in model.state_dict().items():
            if v.dtype.is_floating_point:
                ema[k].mul_(decay).add_(v, alpha=1 - decay)
            else:
                ema[k].copy_(v)


# ─── Source Embedding ───

class SourceEmbedding(nn.Module):
    def __init__(self, num_sources, d_model):
        super().__init__()
        self.embedding = nn.Embedding(num_sources, d_model)
        self.linear = nn.Linear(d_model, d_model)

    def forward(self, x):
        return self.linear(self.embedding(x))


# ─── Data Generation ───

def generate_random_data(start_date, end_date, num_samples, time_scale='s'):
    start = np.datetime64(start_date)
    end = np.datetime64(end_date)
    start_int = start.astype(f'datetime64[{time_scale}]').astype(int)
    end_int = end.astype(f'datetime64[{time_scale}]').astype(int)
    dates = np.linspace(start_int, end_int, num_samples).astype(f'datetime64[{time_scale}]')
    values = np.random.randn(num_samples)
    return dates, values


class SourcePredictionGenerator:
    """Generate source predictions with sinusoidal expertise cycles.

    Each source's expertise follows sin(2π(day + phase) / cycle).
    Exactly one expert per timestep (argmax). Expert source outputs
    the true value; non-experts output random noise.
    """
    def __init__(self, true_data, num_sources, dates, expertise_cycle=365):
        self.true_data = true_data
        self.num_sources = num_sources
        self.dates = dates
        self.expertise_cycle = expertise_cycle
        self.time_vals = (self.dates - self.dates[0]).astype('timedelta64[D]').astype(float)
        print(f"SourcePredictionGenerator: {self.dates[0]} to {self.dates[-1]}")

    def generate_predictions(self):
        predictions = np.zeros((self.num_sources, *self.true_data.shape))
        phases = np.arange(self.num_sources) * (self.expertise_cycle / self.num_sources)
        sine_values = np.sin(2 * np.pi * (self.time_vals[np.newaxis, :] + phases[:, np.newaxis]) / self.expertise_cycle)
        expert_idx = np.argmax(sine_values, axis=0)
        self.expert_mask = np.zeros_like(sine_values, dtype=bool)
        self.expert_mask[expert_idx, np.arange(sine_values.shape[1])] = True

        for source in range(self.num_sources):
            random_preds = np.random.randn(*self.true_data.shape)
            mask = self.expert_mask[source]
            predictions[source] = np.where(mask[np.newaxis, :, np.newaxis], self.true_data, random_preds)

        return predictions


# ─── Multi-Scale Sampling ───

def sample_multiscale_indices(num_datapoints, K, num_samples):
    """Sample K timestamps per example at varying time scales.

    Each sample picks a random span size (log-uniform from small to full range),
    then places K timestamps evenly within that span. Different samples get
    different span sizes so the model sees patterns at every scale.
    """
    all_indices = []
    for _ in range(num_samples):
        min_span = K
        max_span = num_datapoints - 1
        span = int(np.exp(np.random.uniform(np.log(min_span), np.log(max_span))))
        span = min(span, max_span)
        start = np.random.randint(0, num_datapoints - span)
        indices = np.linspace(start, start + span, K, dtype=int)
        indices = np.clip(indices, 0, num_datapoints - 1)
        all_indices.append(indices)
    return np.array(all_indices)


class MultiScaleDataset(Dataset):
    """Each sample: K timestamps, each with N source predictions and a target."""

    def __init__(self, source_predictions, true_data, dates, num_samples, K):
        self.num_sources = source_predictions.shape[0]
        self.T = true_data.shape[1]
        self.K = K
        self.indices = sample_multiscale_indices(self.T, K, num_samples)
        self.all_preds = source_predictions[:, 0, :, 0]  # (N, T)
        self.all_targets = true_data[0, :, 0]  # (T,)
        self.all_timestamps = np.array(
            [dates[i].astype('datetime64[s]').astype(np.int64) for i in range(self.T)],
            dtype=np.int64
        )
        self.source_ids = np.arange(self.num_sources)

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, idx):
        k_indices = self.indices[idx]
        x_preds = self.all_preds[:, k_indices].T  # (K, N)
        targets = self.all_targets[k_indices]  # (K,)
        timestamps = self.all_timestamps[k_indices]  # (K,)
        source_ids = np.tile(self.source_ids, (self.K, 1))  # (K, N)
        return (
            torch.FloatTensor(x_preds),
            torch.FloatTensor(timestamps),
            torch.LongTensor(source_ids),
            torch.FloatTensor(targets),
        )


# ─── Model ───

class TrustTransformerV10(nn.Module):
    """Multi-scale trust transformer.

    Input: K timestamps x N sources = K*N tokens.
    Each token: cat([v, t, s, t*s]), dedicated subspaces.
    Self-attention across all K*N tokens.
    Cross-attention: K queries attend to K*N context.
    Output: K predictions.
    """

    def __init__(self, d_model, nhead, num_layers, num_sources,
                 ts_min=0, ts_max=1, use_temporal=True, use_source=True):
        super().__init__()
        self.d_model = d_model
        self.num_sources = num_sources
        self.use_temporal = use_temporal
        self.use_source = use_source
        self.value_embedding = nn.Linear(1, d_model)

        n_parts = 1  # v always
        if use_temporal: n_parts += 1
        if use_source: n_parts += 1
        if use_temporal and use_source: n_parts += 1
        self.d_model_total = d_model * n_parts

        self.timestamp_enc = NormalizedTimestampEncoding(self.d_model_total, ts_min, ts_max)
        self.context_proj = nn.Linear(self.d_model_total, d_model)
        self.source_embedding = SourceEmbedding(num_sources, d_model)

        self.self_attn_layers = nn.ModuleList([
            TransformerLayer(self.d_model_total, nhead, self.d_model_total * 4)
            for _ in range(num_layers)
        ])

        self.cross_attn = nn.MultiheadAttention(
            embed_dim=self.d_model_total, num_heads=nhead, batch_first=True
        )

        self.fc = nn.Linear(self.d_model_total, 1)

    def forward(self, x_preds, timestamps, source_ids):
        B, K, N = x_preds.shape

        x_flat = x_preds.reshape(B, K * N)
        v = self.value_embedding(x_flat.unsqueeze(-1))

        parts = [v]

        if self.use_temporal:
            t_full = self.timestamp_enc(timestamps)
            t_proj = self.context_proj(t_full)
            t_per_token = t_proj.unsqueeze(2).expand(-1, -1, N, -1).reshape(B, K * N, self.d_model)
            parts.append(t_per_token)

        if self.use_source:
            s_flat = source_ids.reshape(B, K * N)
            s = self.source_embedding(s_flat)
            parts.append(s)

        if self.use_temporal and self.use_source:
            parts.append(t_per_token * s)

        x = torch.cat(parts, dim=-1)

        for layer in self.self_attn_layers:
            x = layer(x)

        # Cross-attention with timestamp query, or mean pooling if no temporal
        if self.use_temporal:
            query = t_full
            x_cross, _ = self.cross_attn(query, x, x)
        else:
            # No temporal info: pool over source tokens per timestamp
            x_reshaped = x.reshape(B, K, N, self.d_model_total)
            x_cross = x_reshaped.mean(dim=2)  # (B, K, d_model_total)

        return self.fc(x_cross).squeeze(-1)


# ─── Evaluation ───

def compute_accuracy(model, loader, num_sources):
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for batch in loader:
            x_preds, timestamps, source_ids, targets = [b.to(device) for b in batch]
            outputs = model(x_preds, timestamps, source_ids)
            out_np = outputs.cpu().numpy()
            tgt_np = targets.cpu().numpy()
            xp_np = x_preds.cpu().numpy()
            B, K, N = xp_np.shape
            for b in range(B):
                for k in range(K):
                    expert = np.argmin(np.abs(xp_np[b, k] - tgt_np[b, k]))
                    pick = np.argmin(np.abs(xp_np[b, k] - out_np[b, k]))
                    if pick == expert:
                        correct += 1
                    total += 1
    return correct / total if total > 0 else 0


def evaluate_loss(model, loader, criterion):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for batch in loader:
            x_preds, timestamps, source_ids, targets = [b.to(device) for b in batch]
            outputs = model(x_preds, timestamps, source_ids)
            total_loss += criterion(outputs, targets).item()
    return total_loss / len(loader)


def collect_predictions(model, loader):
    model.eval()
    all_preds, all_targets, all_x_preds, all_timestamps = [], [], [], []
    with torch.no_grad():
        for batch in loader:
            x_preds, timestamps, source_ids, targets = [b.to(device) for b in batch]
            outputs = model(x_preds, timestamps, source_ids)
            B, K, N = x_preds.shape
            for b in range(B):
                for k in range(K):
                    all_preds.append(outputs[b, k].cpu().numpy())
                    all_targets.append(targets[b, k].cpu().numpy())
                    all_x_preds.append(x_preds[b, k].cpu().numpy())
                    all_timestamps.append(timestamps[b, k].cpu().numpy())
    return {
        'preds': np.array(all_preds),
        'targets': np.array(all_targets),
        'x_preds': np.array(all_x_preds),
        'timestamps': np.array(all_timestamps),
    }


# ─── Diagnostics ───

def plot_confusion_matrix(collected, label, num_sources):
    preds = collected['preds']
    targets = collected['targets']
    x_preds = collected['x_preds']

    true_experts = np.array([np.argmin(np.abs(x_preds[i] - targets[i])) for i in range(len(targets))])
    model_picks = np.array([np.argmin(np.abs(x_preds[i] - preds[i])) for i in range(len(targets))])

    cm = np.zeros((num_sources, num_sources), dtype=int)
    for te, mp in zip(true_experts, model_picks):
        cm[te, mp] += 1

    cm_norm = cm.astype(float) / (cm.sum(axis=1, keepdims=True) + 1e-9)

    fig, axes = plt.subplots(1, 2, figsize=(14, 6))

    ax = axes[0]
    im = ax.imshow(cm, cmap='Blues')
    ax.set_xlabel('Model Picked')
    ax.set_ylabel('True Expert')
    ax.set_title(f'{label}: Confusion (counts)')
    ax.set_xticks(range(num_sources))
    ax.set_yticks(range(num_sources))
    for i in range(num_sources):
        for j in range(num_sources):
            ax.text(j, i, str(cm[i, j]), ha='center', va='center',
                    color='white' if cm[i, j] > cm.max() * 0.5 else 'black', fontsize=8)
    plt.colorbar(im, ax=ax)

    ax = axes[1]
    im = ax.imshow(cm_norm, cmap='Blues', vmin=0, vmax=1)
    ax.set_xlabel('Model Picked')
    ax.set_ylabel('True Expert')
    ax.set_title(f'{label}: Confusion (normalized)')
    ax.set_xticks(range(num_sources))
    ax.set_yticks(range(num_sources))
    for i in range(num_sources):
        for j in range(num_sources):
            ax.text(j, i, f'{cm_norm[i, j]:.2f}', ha='center', va='center',
                    color='white' if cm_norm[i, j] > 0.5 else 'black', fontsize=8)
    plt.colorbar(im, ax=ax)

    per_source_acc = np.diag(cm) / (cm.sum(axis=1) + 1e-9)
    acc_str = ', '.join(f'S{i}: {a:.0%}' for i, a in enumerate(per_source_acc))
    plt.suptitle(f'{label}: Per-source: {acc_str}', fontsize=11, fontweight='bold')
    plt.tight_layout()
    tag = label.lower().replace(' ', '_')
    path = os.path.join(PLOT_DIR, f'v10_{tag}_confusion.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  Saved: {path}")


def plot_expert_cycle(collected, dates, generator, label, num_sources, expertise_cycle):
    timestamps = collected['timestamps']
    targets = collected['targets']
    preds = collected['preds']
    x_preds = collected['x_preds']

    true_experts = np.array([np.argmin(np.abs(x_preds[i] - targets[i])) for i in range(len(targets))])
    model_picks = np.array([np.argmin(np.abs(x_preds[i] - preds[i])) for i in range(len(targets))])

    ts_min = timestamps.min()
    days = (timestamps - ts_min) / 86400.0

    sort_idx = np.argsort(days)
    days = days[sort_idx]
    true_experts = true_experts[sort_idx]
    model_picks = model_picks[sort_idx]

    num_bins = 100
    bin_edges = np.linspace(days.min(), days.max(), num_bins + 1)
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2

    bin_accs, bin_true_expert, bin_model_pick = [], [], []
    for b in range(num_bins):
        mask = (days >= bin_edges[b]) & (days < bin_edges[b + 1])
        if mask.sum() == 0:
            bin_accs.append(np.nan)
            bin_true_expert.append(-1)
            bin_model_pick.append(-1)
            continue
        bin_accs.append(np.mean(true_experts[mask] == model_picks[mask]))
        bin_true_expert.append(np.argmax(np.bincount(true_experts[mask], minlength=num_sources)))
        bin_model_pick.append(np.argmax(np.bincount(model_picks[mask], minlength=num_sources)))

    bin_accs = np.array(bin_accs)
    bin_true_expert = np.array(bin_true_expert)

    test_days = np.linspace(days.min(), days.max(), 500)
    abs_test_days = test_days + (timestamps.min() - float(dates[0].astype('datetime64[s]').astype(np.int64))) / 86400.0
    phases = np.arange(num_sources) * (expertise_cycle / num_sources)
    sine_curves = np.sin(2 * np.pi * (abs_test_days[np.newaxis, :] + phases[:, np.newaxis]) / expertise_cycle)
    ideal_expert = np.argmax(sine_curves, axis=0)

    fig, axes = plt.subplots(3, 1, figsize=(16, 12), sharex=True)
    colors = plt.cm.tab10(np.arange(num_sources))

    ax = axes[0]
    for s in range(num_sources):
        ax.plot(test_days, sine_curves[s], color=colors[s], alpha=0.5, linewidth=1, label=f'Source {s}')
        expert_mask = ideal_expert == s
        ax.fill_between(test_days, -1, 1, where=expert_mask, color=colors[s], alpha=0.15)
    ax.set_ylabel('Sine Value')
    ax.set_title(f'{label}: Expert Cycle (cycle={expertise_cycle} days)')
    ax.legend(loc='upper right', fontsize=8, ncol=min(num_sources, 8))
    ax.grid(True, alpha=0.3)

    ax = axes[1]
    bar_width = (bin_edges[1] - bin_edges[0]) * 0.4
    for b in range(num_bins):
        if bin_true_expert[b] >= 0:
            ax.bar(bin_centers[b] - bar_width/2, bin_true_expert[b] + 0.5, bar_width,
                   color=colors[bin_true_expert[b]], alpha=0.6, edgecolor='black', linewidth=0.5)
        if bin_model_pick[b] >= 0:
            mp = np.array(bin_model_pick)[b]
            ax.bar(bin_centers[b] + bar_width/2, mp + 0.5, bar_width,
                   color=colors[mp], alpha=0.3, edgecolor='gray', linewidth=0.5, hatch='//')
    ax.set_ylabel('Source ID')
    ax.set_title('Dominant Expert (solid) vs Model Pick (hatched) per time bin')
    ax.set_yticks(range(num_sources))
    ax.grid(True, alpha=0.3)

    ax = axes[2]
    ax.plot(bin_centers, bin_accs, color='black', linewidth=2)
    ax.axhline(y=1/num_sources, color='red', linestyle='--', alpha=0.5, label=f'Random (1/{num_sources})')
    ax.fill_between(bin_centers, bin_accs, alpha=0.2, color='green')
    ax.set_xlabel('Days into test window')
    ax.set_ylabel('Accuracy')
    ax.set_title(f'Accuracy over time (overall: {np.nanmean(bin_accs):.1%})')
    ax.set_ylim(0, 1)
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.suptitle(f'{label}: Expert Cycle vs Predictions (seed={SEED})', fontsize=13, fontweight='bold')
    plt.tight_layout()
    tag = label.lower().replace(' ', '_')
    path = os.path.join(PLOT_DIR, f'v10_{tag}_cycle.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  Saved: {path}")


def plot_training_curves(train_accs, test_accs, train_losses, test_losses,
                         config_name, num_sources, best_test_acc):
    from scipy import stats as sp_stats
    random_bl = 1.0 / num_sources
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))

    ax = axes[0]
    ax.plot(train_accs, label='Train (EMA)', color='blue', alpha=0.7)
    ax.plot(test_accs, label='Test (EMA)', color='orange', alpha=0.7)
    ax.axhline(y=random_bl, color='red', linestyle='--', alpha=0.5, label=f'Random (1/{num_sources})')

    # Trend analysis on last 50 epochs (or all if fewer)
    trend_window = min(50, len(test_accs))
    if trend_window >= 10:
        tail = np.array(test_accs[-trend_window:])
        x = np.arange(trend_window)
        slope, intercept, r, p, se = sp_stats.linregress(x, tail)
        x_plot = np.arange(len(test_accs) - trend_window, len(test_accs))
        ax.plot(x_plot, intercept + slope * x, color='green', linewidth=2,
                linestyle='-', alpha=0.9, label='Trend (last 50)')
        sig = '***' if p < 0.001 else '**' if p < 0.01 else '*' if p < 0.05 else 'n.s.'
        ax.text(0.02, 0.02,
                f'Slope: {slope*100:+.2f} pp/epoch ({sig}, p={p:.3f})\n'
                f'Projected +{slope*100*50:.1f} pp over next 50 epochs',
                transform=ax.transAxes, fontsize=8, verticalalignment='bottom',
                bbox=dict(boxstyle='round,pad=0.3', facecolor='lightgreen', alpha=0.7))

    ax.set_title(f'{config_name}: Accuracy')
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Accuracy')
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)
    ax.set_ylim(0, 1)

    ax = axes[1]
    ax.plot(train_losses, label='Train (raw)', color='blue', alpha=0.7)
    ax.plot(test_losses, label='Test (EMA)', color='orange', alpha=0.7)
    ax.set_title(f'{config_name}: Loss')
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss (L1)')
    ax.legend()
    ax.grid(True, alpha=0.3)

    plt.suptitle(f'v10 K={K_TIMESTAMPS} EMA={EMA_DECAY} L1: {config_name} (seed={SEED})\n'
                 f'Best test: {best_test_acc:.1%}',
                 fontsize=12, fontweight='bold')
    plt.tight_layout()
    tag = config_name.lower().replace(' ', '_')
    path = os.path.join(PLOT_DIR, f'v10_{tag}_curves.png')
    plt.savefig(path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  Saved: {path}")


# ─── Training ───

def run_config(config_name, num_samples, num_sources, expertise_cycle,
               use_temporal=True, use_source=True, diagnostics=True, num_epochs=None):
    if num_epochs is None:
        num_epochs = NUM_EPOCHS
    torch.manual_seed(SEED)
    np.random.seed(SEED)

    print(f"\n{'='*70}")
    print(f"v10: {config_name}")
    print(f"  {num_samples} samples, K={K_TIMESTAMPS}, {num_epochs} epochs, ema_decay={EMA_DECAY}")
    print(f"  {num_sources} sources, cycle={expertise_cycle}")
    print(f"  use_temporal={use_temporal}, use_source={use_source}")
    print(f"{'='*70}")

    dates, values = generate_random_data('2015-01-01', '2022-09-27', num_samples)
    true_data = values.reshape(1, -1, 1)
    generator = SourcePredictionGenerator(true_data, num_sources, dates,
                                          expertise_cycle=expertise_cycle)
    source_predictions = generator.generate_predictions()

    dataset = MultiScaleDataset(source_predictions, true_data, dates, num_samples, K_TIMESTAMPS)

    ts_min = float(dataset.all_timestamps.min())
    ts_max = float(dataset.all_timestamps.max())

    # Chronological split on first timestamp of each sample
    split_time = ts_min + 0.8 * (ts_max - ts_min)
    train_idx = [i for i in range(len(dataset))
                 if dataset.all_timestamps[dataset.indices[i, 0]] < split_time]
    test_idx = [i for i in range(len(dataset))
                if dataset.all_timestamps[dataset.indices[i, 0]] >= split_time]
    train_ds = torch.utils.data.Subset(dataset, train_idx)
    test_ds = torch.utils.data.Subset(dataset, test_idx)

    print(f"  Train: {len(train_ds)}, Test: {len(test_ds)}")

    samples_per_epoch = BATCH_PER_EPOCH * BATCH_SIZE
    train_sampler = RandomSampler(train_ds, replacement=True, num_samples=samples_per_epoch)
    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, sampler=train_sampler)
    test_loader = DataLoader(test_ds, batch_size=BATCH_SIZE)

    model = TrustTransformerV10(
        d_model=32, nhead=4, num_layers=2, num_sources=num_sources,
        ts_min=ts_min, ts_max=ts_max,
        use_temporal=use_temporal, use_source=use_source
    ).to(device)

    optimizer = torch.optim.Adam(model.parameters())
    criterion = nn.L1Loss()   # THE one change from trust_v9

    ema = clone_state(model)

    train_accs, test_accs = [], []
    train_losses, test_losses = [], []
    best_test_acc = 0.0
    best_state = None

    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        n_batches = 0
        for batch_data in train_loader:
            x_preds, timestamps, source_ids, targets = [b.to(device) for b in batch_data]

            # Value dropout
            mask = (torch.rand_like(x_preds) > 0.5).float()
            x_preds_masked = x_preds * mask

            optimizer.zero_grad()
            out = model(x_preds_masked, timestamps, source_ids)
            loss = criterion(out, targets)
            loss.backward()
            optimizer.step()
            ema_update(model, ema, EMA_DECAY)
            total_loss += loss.item()
            n_batches += 1

        train_loss = total_loss / n_batches

        # Evaluate the EMA weights (the v9 method)
        backup = clone_state(model)
        model.load_state_dict(ema)
        test_loss = evaluate_loss(model, test_loader, criterion)
        train_acc = compute_accuracy(model, train_loader, num_sources)
        test_acc = compute_accuracy(model, test_loader, num_sources)
        model.load_state_dict(backup)

        if test_acc > best_test_acc:
            best_test_acc = test_acc
            best_state = {k: v.clone() for k, v in ema.items()}

        train_accs.append(train_acc)
        test_accs.append(test_acc)
        train_losses.append(train_loss)
        test_losses.append(test_loss)

        if (epoch + 1) % 10 == 0:
            print(f"  Epoch {epoch+1}: train={train_acc:.3f} test={test_acc:.3f} (best={best_test_acc:.3f})", flush=True)

    if best_state is not None:
        model.load_state_dict(best_state)

    tag = config_name.lower().replace(' ', '_').replace('(', '').replace(')', '')
    ckpt_path = os.path.join(CKPT_DIR, f'v10_{tag}_seed{SEED}.pt')
    torch.save({
        'state_dict': best_state, 'best_test_acc': best_test_acc,
        'K': K_TIMESTAMPS, 'ema_decay': EMA_DECAY, 'seed': SEED, 'loss': 'L1',
        'num_sources': num_sources, 'expertise_cycle': expertise_cycle,
        'ts_min': ts_min, 'ts_max': ts_max,
        'use_temporal': use_temporal, 'use_source': use_source,
    }, ckpt_path)
    print(f"  Checkpoint: {ckpt_path}")

    if diagnostics:
        print(f"\n--- Parameter Inspection ---")
        inspect_encoding(model.timestamp_enc, f"v10 {config_name}")
        plot_training_curves(train_accs, test_accs, train_losses, test_losses,
                             config_name, num_sources, best_test_acc)
        collected = collect_predictions(model, test_loader)
        plot_confusion_matrix(collected, config_name, num_sources)
        plot_expert_cycle(collected, dates, generator, config_name, num_sources, expertise_cycle)

    return best_test_acc


# ─── Main ───

if __name__ == '__main__':
    print(f"Trust Proof v10: Multi-Scale Context Window + EMA Evaluation + L1 Loss")
    print(f"  K={K_TIMESTAMPS}, full={NUM_EPOCHS} epochs, ablation={NUM_EPOCHS_ABLATION} epochs, seed={SEED}")
    print(f"  {NUM_SAMPLES} samples, {BATCH_PER_EPOCH} batches/epoch, batch_size={BATCH_SIZE}, ema_decay={EMA_DECAY}")
    print(f"  Device: {device}")

    run_hard = '--hard' in sys.argv or ('--simple' not in sys.argv)
    run_simple = '--simple' in sys.argv or ('--hard' not in sys.argv)
    full_only = '--full-only' in sys.argv

    results = {}

    # Full models with diagnostics (200 epochs)
    if run_hard:
        results['HARD'] = run_config("HARD", NUM_SAMPLES, 8, 365)
    if run_simple:
        results['SIMPLE'] = run_config("SIMPLE", NUM_SAMPLES, 3, 500)

    # Ablations (50 epochs)
    if not full_only:
        ablations = [
            ("Baseline (Value Only)", False, False),
            ("Temporal Only", True, False),
            ("Source Only", False, True),
        ]
        for abl_name, use_t, use_s in ablations:
            if run_hard:
                label = f"HARD {abl_name}"
                results[label] = run_config(label, NUM_SAMPLES, 8, 365,
                                            use_temporal=use_t, use_source=use_s,
                                            diagnostics=False, num_epochs=NUM_EPOCHS_ABLATION)
            if run_simple:
                label = f"SIMPLE {abl_name}"
                results[label] = run_config(label, NUM_SAMPLES, 3, 500,
                                            use_temporal=use_t, use_source=use_s,
                                            diagnostics=False, num_epochs=NUM_EPOCHS_ABLATION)

    # Summary
    print(f"\n{'='*70}")
    print(f"TRUST PROOF v10 SUMMARY (K={K_TIMESTAMPS}, EMA={EMA_DECAY}, L1, seed={SEED})")
    print(f"{'='*70}")
    for label, acc in results.items():
        print(f"  {label}: {acc:.3f}")