"""
Dataset Distancing Tutorial
============================
End-to-end: generate datasets from DeepMIMO, split into train/val/test,
compute distances and transfer performance, measure correlation.

Usage:
    python tutorial.py
    python tutorial.py --trials 5
    python tutorial.py --task beam-pred --n_beams 8 --epochs 150 --batch_size 512
"""

import argparse, logging, sys, os

parser = argparse.ArgumentParser(description="Dataset Distancing Tutorial")
# paths
parser.add_argument("--lwm_path",      type=str, default="third_party/lwm-1.1")
parser.add_argument("--weights_file",  type=str, default=None, help="alternative checkpoint (e.g. CDE-finetuned)")
parser.add_argument("--cache_dir",     type=str, default="cache/tutorial")
parser.add_argument("--gpu",           type=int, default=1)
# data splits
parser.add_argument("--n_train", type=int, default=3000)
parser.add_argument("--n_val",   type=int, default=2000)
parser.add_argument("--n_test",  type=int, default=2000)
parser.add_argument("--seed",    type=int, default=42)
parser.add_argument("--trials",  type=int, default=1, help="number of trials with different split seeds")
# scenario generation
parser.add_argument("--task",    type=str, default="los-nlos",
                    choices=["los-nlos", "beam-pred"])
parser.add_argument("--n_ant",   type=int, default=32)
parser.add_argument("--n_sub",   type=int, default=32)
parser.add_argument("--n_beams", type=int, default=32, help="codebook size for beam-pred")
# distance computation
parser.add_argument("--distance",       type=str, default="swd", choices=["swd", "centroid", "cosine"])
parser.add_argument("--embed_space",    type=str, default="lwm", choices=["lwm", "raw", "umap"],
                    help="embedding space for distance computation")
parser.add_argument("--umap_dim",       type=int, default=2, help="UMAP target dimension")
parser.add_argument("--label_aware",    action="store_true", default=True)
parser.add_argument("--no_label_aware", action="store_true")
parser.add_argument("--sw_projections", type=int, default=64)
parser.add_argument("--asymmetric",    action="store_true", help="directed SWD + unsymmetrized P")
parser.add_argument("--raw_perf",      action="store_true", help="skip P symmetrization (keep symmetric D)")
parser.add_argument("--asym_alpha",    type=float, default=1.0, help="strength of asymmetry")
# CDE fine-tuning
parser.add_argument("--cde",          action="store_true", help="train CDE before computing distances")
parser.add_argument("--cde_epochs",   type=int, default=40)
parser.add_argument("--cde_lr",       type=float, default=5e-4)
parser.add_argument("--cde_save",     type=str, default=None, help="save CDE checkpoint to this dir")
# performance matrix
parser.add_argument("--epochs",     type=int, default=15)
parser.add_argument("--lr",         type=float, default=1e-3)
parser.add_argument("--batch_size", type=int, default=128)
# label robustness
parser.add_argument("--label_noise", type=float, default=0.0, help="fraction of labels to flip randomly")
parser.add_argument("--label_drop",  type=float, default=0.0, help="fraction of labels to drop (→ label-agnostic D)")
parser.add_argument("--log_dir",     type=str, default="logs", help="directory for log files")
args = parser.parse_args()

if args.no_label_aware:
    args.label_aware = False

os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

# ── Logging setup ────────────────────────────────────────────────────────────
os.makedirs(args.log_dir, exist_ok=True)
from datetime import datetime
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_file = os.path.join(args.log_dir, f"{args.task}_{timestamp}.log")

log = logging.getLogger("tutorial")
log.setLevel(logging.INFO)
log.handlers.clear()
# File handler
fh = logging.FileHandler(log_file)
fh.setFormatter(logging.Formatter("%(asctime)s %(message)s", datefmt="%H:%M:%S"))
log.addHandler(fh)
# Console handler
ch = logging.StreamHandler()
ch.setFormatter(logging.Formatter("%(message)s"))
log.addHandler(ch)

log.info(f"Command: {' '.join(sys.argv)}")
log.info(f"Args: {vars(args)}")
log.info(f"Log file: {log_file}")

import torch, numpy as np

from dataset_distancing import (
    DatasetCollection,
    compute_distances,
    compute_performance_matrix,
    correlate,
    DistanceConfig,
    TransferConfig,
)
from dataset_distancing.embeddings import register_lwm_path, load_runner

register_lwm_path(args.lwm_path)

# ── Step 1: Generate datasets (once, shared across trials) ───────────────────
import deepmimo as dm
scenarios = (
    [{"name": s, "tx_set": tx}
     for s in dm.search()[1:5]
     for tx in [1]]
    + dm.search()[60:65]
)
full = DatasetCollection.from_scenarios(
    scenarios,
    task=args.task,
    n_ant=args.n_ant,
    n_sub=args.n_sub,
    n_beams=args.n_beams,
    cache_dir=args.cache_dir,
)

def _print_matrix(name, M):
    log.info(f"\n{name}:")
    for row in M:
        log.info("  " + "  ".join(f"{v:.4f}" for v in row))

def _perturb_labels(collection, noise_frac, drop_frac, seed):
    """Flip and/or drop a fraction of labels in each dataset."""
    from dataset_distancing.data import WirelessDataset
    g = torch.Generator().manual_seed(seed)
    new_ds = []
    for ds in collection.datasets:
        if ds.labels is None:
            new_ds.append(ds)
            continue
        labels = ds.labels.clone()
        n = len(labels)
        # Flip labels
        if noise_frac > 0:
            n_flip = int(noise_frac * n)
            flip_idx = torch.randperm(n, generator=g)[:n_flip]
            n_classes = int(labels.max().item()) + 1
            labels[flip_idx] = torch.randint(0, n_classes, (n_flip,), generator=g)
        # Drop labels
        if drop_frac > 0:
            n_drop = int(drop_frac * n)
            drop_idx = torch.randperm(n, generator=g)[:n_drop]
            labels[drop_idx] = -1  # sentinel
            # Keep only samples with valid labels
            keep = labels >= 0
            new_ds.append(WirelessDataset(
                name=ds.name,
                channels=ds.channels[keep],
                labels=labels[keep],
                targets=ds.targets[keep] if ds.targets is not None else None,
            ))
            continue
        new_ds.append(WirelessDataset(
            name=ds.name, channels=ds.channels, labels=labels,
            targets=ds.targets,
        ))
    return DatasetCollection(datasets=new_ds)

# ── Run trials ───────────────────────────────────────────────────────────────
all_pearson, all_spearman = [], []

for trial in range(args.trials):
    split_seed = args.seed + trial

    torch.manual_seed(split_seed)
    np.random.seed(split_seed)

    # ── Split ────────────────────────────────────────────────────────────
    train_ds, val_ds, test_ds = [], [], []
    for ds in full:
        tr, va, te = ds.split(
            train=args.n_train, val=args.n_val, test=args.n_test, seed=split_seed,
        )
        train_ds.append(tr)
        val_ds.append(va)
        test_ds.append(te)
    train_col = DatasetCollection(datasets=train_ds)
    val_col   = DatasetCollection(datasets=val_ds)
    test_col  = DatasetCollection(datasets=test_ds)

    # ── Perturb labels for distance computation only ──────────────────
    # Keep clean labels for P (ground-truth transfer performance),
    # perturb only the copies used for label-aware D.
    if args.label_noise > 0 or args.label_drop > 0:
        train_col_noisy = _perturb_labels(train_col, args.label_noise, args.label_drop, split_seed)
    else:
        train_col_noisy = train_col

    if trial == 0:
        log.info(f"\n{'='*60}")
        log.info(f"Generated {len(full)} datasets "
              f"(train/val/test = {args.n_train}/{args.n_val}/{args.n_test}):")
        for tr, te in zip(train_col, test_col):
            n0 = (tr.labels == 0).sum().item()
            log.info(f"  {tr.name}: train={len(tr)}, test={len(te)}, "
                  f"LoS%={100*n0/max(len(tr),1):.0f}%")

    # ── Load runner ──────────────────────────────────────────────────────
    import io, contextlib
    with contextlib.redirect_stdout(io.StringIO()):
        runner = load_runner(args.lwm_path, weights_file=args.weights_file)

    # ── P matrix ─────────────────────────────────────────────────────────
    verbose = (args.trials == 1)
    P = compute_performance_matrix(
        train_col,
        cfg=TransferConfig(
            task=args.task,
            epochs=args.epochs,
            lr=args.lr,
            batch_size=args.batch_size,
            n_beams=args.n_beams,
            seed=split_seed,
            symmetrize=not (args.asymmetric or args.raw_perf),
            verbose=verbose,
        ),
        val_collection=val_col,
        test_collection=test_col,
    )

    # ── CDE (optional) ───────────────────────────────────────────────────
    if args.cde:
        from dataset_distancing.cde import CDETrainer
        from dataset_distancing import CDEConfig

        cde_cfg = CDEConfig(
            epochs=args.cde_epochs,
            lr=args.cde_lr,
            n_per_dataset=64,
            n_per_dataset_eval=128,
            corr_n_per_dataset=64,
            asymmetric=args.asymmetric,
            distance=DistanceConfig(
                sw_projections=args.sw_projections,
                label_aware=args.label_aware,
                asymmetric=args.asymmetric,
            ),
        )
        trainer = CDETrainer(runner, cde_cfg)
        trainer.fit(train_col, P, eval_collection=train_col, perf_eval=P, verbose=verbose)
        if args.cde_save and trial == 0:
            trainer.save(args.cde_save)

    # ── D matrix ─────────────────────────────────────────────────────────
    if args.embed_space == "lwm":
        D = compute_distances(train_col_noisy, runner, cfg=DistanceConfig(
            distance=args.distance,
            label_aware=args.label_aware,
            n_per_dataset=len(train_col[0]),
            sw_projections=args.sw_projections,
            asymmetric=args.asymmetric,
            asymmetric_alpha=args.asym_alpha,
        ))
    else:
        # Raw or UMAP: build EmbeddingResult manually
        from dataset_distancing.embeddings import EmbeddingResult, channels_to_patches
        from dataset_distancing.distances import compute_distances_from_embeddings

        # Vectorize channels: (N, 1, ant, sub) complex → (N, 2*ant*sub) float
        # Subsample to common size across datasets
        n_per = min(len(ds) for ds in train_col_noisy.datasets)
        all_embs, all_labels = [], []
        for ds in train_col_noisy.datasets:
            idx = torch.randperm(len(ds))[:n_per]
            ch = ds.channels[idx]
            labels = ds.labels[idx] if ds.labels is not None else None

            if ch.is_complex():
                flat = torch.cat([ch.real.reshape(n_per, -1),
                                  ch.imag.reshape(n_per, -1)], dim=1).float()
            else:
                flat = ch.reshape(n_per, -1).float()

            if args.embed_space == "umap":
                import warnings, umap
                from sklearn.preprocessing import StandardScaler
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    scaled = StandardScaler().fit_transform(flat.numpy())
                    reducer = umap.UMAP(n_components=args.umap_dim, n_neighbors=16,
                                        min_dist=0.1, random_state=split_seed)
                    flat = torch.from_numpy(reducer.fit_transform(scaled)).float()

            all_embs.append(flat.unsqueeze(0))
            all_labels.append(labels)

        emb_result = EmbeddingResult(
            embeddings=torch.cat(all_embs, dim=0),
            labels=all_labels,
            dataset_names=train_col_noisy.names(),
        )
        D = compute_distances_from_embeddings(emb_result, cfg=DistanceConfig(
            distance=args.distance,
            label_aware=args.label_aware,
            sw_projections=args.sw_projections,
            asymmetric=args.asymmetric,
            asymmetric_alpha=args.asym_alpha,
        ))

    # ── Correlate ────────────────────────────────────────────────────────
    stats = correlate(D, P, symmetric=not (args.asymmetric or args.raw_perf))
    all_pearson.append(stats['pearson'])
    all_spearman.append(stats['spearman'])

    if args.trials == 1:
        _print_matrix("Distance matrix D", D.numpy())
        _print_matrix("Performance matrix P", P.numpy())
        log.info(f"\n{'='*60}")
        log.info(f"Pearson  = {stats['pearson']:.4f}")
        log.info(f"Spearman = {stats['spearman']:.4f}")
        log.info(f"{'='*60}")
    else:
        log.info(f"  Trial {trial+1}/{args.trials} (seed={split_seed}): "
              f"r={stats['pearson']:.4f}, ρ={stats['spearman']:.4f}")

# ── Summary ──────────────────────────────────────────────────────────────────
if args.trials > 1:
    p = np.array(all_pearson)
    s = np.array(all_spearman)
    log.info(f"\n{'='*60}")
    log.info(f"Over {args.trials} trials:")
    log.info(f"  Pearson  = {p.mean():.4f} ± {p.std():.4f}")
    log.info(f"  Spearman = {s.mean():.4f} ± {s.std():.4f}")
    log.info(f"{'='*60}")


# Examples, from the repository root:
#   python tutorial.py --task los-nlos --n_train 400 --n_test 400 --epochs 200 --batch_size 32 --seed 0 --trials 20
#   python tutorial.py --task beam-pred --n_beams 32 --n_train 2000 --n_test 2000 --epochs 200 --batch_size 512 --seed 0 --trials 20
#   python tutorial.py --task los-nlos --embed_space raw       # the same on raw channels
#   python tutorial.py --task los-nlos --cde --cde_save ckpt    # fine-tune with the LWM-CDE objective first