#!/usr/bin/env python3
"""SSL pretraining on LIVECell images.

Methods: SimCLR, SupCon-cellline, MS-CLACL (morphology-soft).
Backbone: ResNet50, exportable to Mask R-CNN.
"""

from __future__ import annotations

import argparse
import json
import random
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torchvision.models import resnet50, ResNet50_Weights
from torchvision.transforms import InterpolationMode


# ---------------------------------------------------------------------------
# augmentations
# ---------------------------------------------------------------------------

class SimCLRTransform:
    """Two-view augmentation for contrastive learning."""

    def __init__(self, size: int = 224, s: float = 1.0):
        self.train_transform = transforms.Compose([
            transforms.RandomResizedCrop(size, scale=(0.2, 1.0), interpolation=InterpolationMode.BICUBIC),
            transforms.RandomApply([transforms.ColorJitter(0.4 * s, 0.4 * s, 0.4 * s, 0.1 * s)], p=0.8),
            transforms.RandomGrayscale(p=0.2),
            transforms.RandomApply([transforms.GaussianBlur(kernel_size=int(0.1 * size) // 2 * 2 + 1, sigma=(0.1, 2.0))], p=0.5),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def __call__(self, x):
        return self.train_transform(x), self.train_transform(x)


# ---------------------------------------------------------------------------
# dataset
# ---------------------------------------------------------------------------

from cell_line_utils import CELL_LINE_NAMES, infer_cell_line


class LiveCellImageDataset(Dataset):
    """Load all LIVECell train images without annotations."""

    def __init__(
        self,
        data_root: Path,
        ann_file: Path,
        transform: Any,
        cell_line_labels: bool = False,
    ):
        self.data_root = Path(data_root)
        self.transform = transform
        self.need_labels = cell_line_labels

        with open(ann_file, "r") as f:
            coco = json.load(f)

        self.image_paths: List[Path] = []
        self.image_ids: List[int] = []
        self.cell_lines: List[int] = []

        for img_info in coco.get("images", []):
            file_name = img_info.get("file_name", "")
            if not file_name:
                continue
            path = self._find_image(file_name)
            if path is None:
                continue
            self.image_paths.append(path)
            self.image_ids.append(int(img_info.get("id", -1)))
            if cell_line_labels:
                self.cell_lines.append(infer_cell_line(file_name))

        if not self.image_paths:
            raise RuntimeError(f"No images found for {ann_file}")

    def _find_image(self, file_name: str) -> Optional[Path]:
        stem = Path(file_name).stem
        for ext in [".tif", ".tiff", ".png", ".jpg", ".jpeg", ".bmp"]:
            # Direct path
            direct = self.data_root / (stem + ext)
            if direct.exists():
                return direct
            # Under images/ subdirs
            for sub in ["images/livecell_train_val_images", "images/livecell_test_images", "images"]:
                candidate = self.data_root / sub / (stem + ext)
                if candidate.exists():
                    return candidate
        # Fallback: search all subdirs under images/
        images_dir = self.data_root / "images"
        if images_dir.exists():
            for p in images_dir.rglob(stem + ".*"):
                return p
        return None

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, index: int) -> Tuple[Any, ...]:
        with Image.open(str(self.image_paths[index])) as img:
            img = img.convert("RGB")
        views = self.transform(img)  # (view1, view2) or single view
        if self.need_labels:
            return views[0], views[1], self.cell_lines[index]
        return views[0], views[1]


# ---------------------------------------------------------------------------
# SSL losses
# ---------------------------------------------------------------------------

def info_nce_loss(features: torch.Tensor, temperature: float = 0.07) -> torch.Tensor:
    """NT-Xent / InfoNCE loss for SimCLR.

    Args:
        features: (2*B, D) where first B are view1, last B are view2.
    """
    B = features.shape[0] // 2
    features = F.normalize(features, dim=1)
    sim = torch.mm(features, features.T) / temperature  # (2B, 2B)

    mask = torch.eye(2 * B, dtype=torch.bool, device=features.device)
    sim = sim.masked_fill(mask, float("-inf"))

    positives = torch.cat([torch.arange(B, 2 * B), torch.arange(0, B)]).to(features.device)
    positives = sim[torch.arange(2 * B), positives].view(2 * B, 1)

    negatives = sim.masked_fill(
        _off_diagonal_mask(2 * B, features.device, include_positives=False),
        float("-inf"),
    )

    logits = torch.cat([positives, negatives], dim=1)
    labels = torch.zeros(2 * B, dtype=torch.long, device=features.device)
    return F.cross_entropy(logits, labels)


def _off_diagonal_mask(n: int, device: torch.device, include_positives: bool = True) -> torch.Tensor:
    """Mask off-diagonal elements of a similarity matrix."""
    B = n // 2
    mask = torch.eye(n, dtype=torch.bool, device=device)
    if not include_positives:
        for i in range(B):
            mask[i, i + B] = True
            mask[i + B, i] = True
    return mask


def supcon_loss(
    features: torch.Tensor,
    labels: torch.Tensor,
    temperature: float = 0.07,
) -> torch.Tensor:
    """Supervised contrastive loss (Khosla et al. 2020, mean-over-positives).

    features: (2*B, D) two views.
    labels: (B,) cell-line labels.
    """
    device = features.device
    features = F.normalize(features, dim=1)

    # replicate labels for both views
    labels = labels.repeat(2)  # (2*B,)
    sim = torch.mm(features, features.T) / temperature  # (2B, 2B)

    # self mask
    self_mask = torch.eye(sim.shape[0], dtype=torch.bool, device=device)
    sim = sim.masked_fill(self_mask, float("-inf"))

    # positive mask: same label, excluding self
    pos_mask = labels.unsqueeze(0) == labels.unsqueeze(1)  # (2B, 2B)
    pos_mask.fill_diagonal_(False)

    # denominator: sum over all except self
    exp_sim = torch.exp(sim)
    denom = exp_sim.sum(dim=1, keepdim=True).clamp(min=1e-9)

    # per-positive log-prob, then mean over positives per anchor.
    # NOTE: sim has -inf on self positions; -inf * 0 == nan in PyTorch.
    # Use torch.where to zero out non-positive positions before summing.
    log_prob = sim - torch.log(denom)
    pos_log_prob = torch.where(
        pos_mask, log_prob, torch.zeros_like(log_prob)
    ).sum(dim=1, keepdim=True)
    num_pos = pos_mask.float().sum(dim=1, keepdim=True).clamp(min=1.0)

    loss = -(pos_log_prob / num_pos).mean()
    return loss


def msclacl_loss(
    features: torch.Tensor,
    image_ids: torch.Tensor,
    id_to_row: Dict[int, int],
    similarity_matrix: torch.Tensor,
    temperature: float = 0.07,
    soft_threshold: float = 0.5,
) -> torch.Tensor:
    """MS-CLACL: morphology-soft contrastive loss.

    Uses precomputed image-level morphology similarity as soft weights.
    Pairs with similarity below soft_threshold are treated as negatives (weight=0).
    Pairs above threshold use their raw similarity as the positive weight.
    The two augmented views of the same image are always a positive pair with
    weight 1.0, preserving the SimCLR base objective.

    features: (2*B, D)
    image_ids: (B,) COCO image IDs of the batch images
    id_to_row: mapping from COCO image_id → row index in similarity_matrix
    similarity_matrix: (N, N) precomputed morphology similarity
    soft_threshold: minimum similarity to be considered as a positive pair
    """
    B = features.shape[0] // 2
    device = features.device
    features = F.normalize(features, dim=1)

    # Map image_ids to matrix row indices
    ids_list = image_ids.cpu().tolist()
    rows = [id_to_row[i] for i in ids_list]
    row_tensor = torch.tensor(rows, dtype=torch.long)

    # get pairwise similarity for the batch
    sim_mat_batch = similarity_matrix[row_tensor][:, row_tensor]  # (B, B)
    # replicate to 2B x 2B (view1+view2 share the same image-level similarity)
    sim_mat_2b = sim_mat_batch.repeat(2, 2).to(device)  # (2B, 2B)

    # compute cosine similarity
    cos_sim = torch.mm(features, features.T) / temperature  # (2B, 2B)

    # self mask
    self_mask = torch.eye(2 * B, dtype=torch.bool, device=device)
    cos_sim = cos_sim.masked_fill(self_mask, float("-inf"))

    # soft weights: above threshold → use raw similarity, below → 0 (negative)
    weights = torch.where(
        sim_mat_2b >= soft_threshold,
        sim_mat_2b,
        torch.zeros_like(sim_mat_2b),
    )
    weights.fill_diagonal_(0.0)  # exclude self
    for i in range(B):
        weights[i, i + B] = 1.0
        weights[i + B, i] = 1.0

    # weighted contrastive loss
    exp_sim = torch.exp(cos_sim) * weights
    pos_sum = exp_sim.sum(dim=1, keepdim=True).clamp(min=1e-9)

    denom = torch.exp(cos_sim).masked_fill(self_mask, 0.0).sum(dim=1, keepdim=True).clamp(min=1e-9)

    loss = -torch.log(pos_sum / denom).mean()
    return loss


# ---------------------------------------------------------------------------
# model
# ---------------------------------------------------------------------------

class SSLModel(nn.Module):
    """ResNet50 backbone + 2-layer projection head."""

    def __init__(self, feature_dim: int = 128, proj_hidden_dim: int = 2048):
        super().__init__()
        self.backbone = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        self.projection = nn.Sequential(
            nn.Linear(2048, proj_hidden_dim),
            nn.BatchNorm1d(proj_hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(proj_hidden_dim, feature_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.backbone.conv1(x)
        h = self.backbone.bn1(h)
        h = self.backbone.relu(h)
        h = self.backbone.maxpool(h)
        h = self.backbone.layer1(h)
        h = self.backbone.layer2(h)
        h = self.backbone.layer3(h)
        h = self.backbone.layer4(h)
        h = self.backbone.avgpool(h)
        h = torch.flatten(h, 1)
        z = self.projection(h)
        return z

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        """Backbone features without projection head (2048-dim)."""
        h = self.backbone.conv1(x)
        h = self.backbone.bn1(h)
        h = self.backbone.relu(h)
        h = self.backbone.maxpool(h)
        h = self.backbone.layer1(h)
        h = self.backbone.layer2(h)
        h = self.backbone.layer3(h)
        h = self.backbone.layer4(h)
        h = self.backbone.avgpool(h)
        return torch.flatten(h, 1)


# ---------------------------------------------------------------------------
# training loop
# ---------------------------------------------------------------------------

def train_ssl(
    model: nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    epochs: int,
    method: str,
    temperature: float,
    similarity_matrix: Optional[torch.Tensor] = None,
    id_to_row: Optional[Dict[int, int]] = None,
    log_interval: int = 50,
) -> List[Dict[str, Any]]:
    """Unified SSL training loop."""
    model.train()
    rows: List[Dict[str, Any]] = []
    global_step = 0

    for epoch in range(epochs):
        epoch_loss = 0.0
        epoch_steps = 0

        for batch in loader:
            started = time.time()

            if method == "simclr":
                x1, x2 = [t.to(device) for t in batch[:2]]
                x = torch.cat([x1, x2], dim=0)  # (2B, C, H, W)
                z = model(x)
                loss = info_nce_loss(z, temperature)

            elif method == "supcon":
                x1, x2, labels = [t.to(device) if isinstance(t, torch.Tensor) else t for t in batch]
                x = torch.cat([x1, x2], dim=0)
                z = model(x)
                loss = supcon_loss(z, labels, temperature)

            elif method == "msclacl":
                x1, x2, ids = [t.to(device) if isinstance(t, torch.Tensor) else t for t in batch]
                x = torch.cat([x1, x2], dim=0)
                z = model(x)
                loss = msclacl_loss(z, ids, id_to_row, similarity_matrix, temperature)

            else:
                raise ValueError(f"Unknown method: {method}")

            if not torch.isfinite(loss):
                raise RuntimeError(f"Non-finite SSL loss at step {global_step}: {loss.item()}")

            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()

            step_loss = loss.item()
            epoch_loss += step_loss
            epoch_steps += 1

            if global_step % log_interval == 0 or global_step == 0:
                print(
                    f"SSL {method} epoch={epoch} step={global_step} loss={step_loss:.4f} "
                    f"time={time.time()-started:.2f}s",
                    flush=True,
                )

            rows.append({"epoch": epoch, "step": global_step, "loss": step_loss})
            global_step += 1

        avg_loss = epoch_loss / max(epoch_steps, 1)
        print(f"SSL {method} epoch={epoch} avg_loss={avg_loss:.4f}", flush=True)

    return rows


# ---------------------------------------------------------------------------
# data loading helpers
# ---------------------------------------------------------------------------

def discover_train_ann(data_root: Path) -> Path:
    candidates = sorted(data_root.rglob("livecell_coco_train.json"))
    if candidates:
        return candidates[0]
    candidates = sorted(data_root.rglob("*train*.json"))
    for c in candidates:
        if "size_split" not in str(c) and "single_cells" not in str(c):
            return c
    raise FileNotFoundError(f"No train COCO JSON found under {data_root}")


def write_csv(path: Path, rows: List[Dict[str, Any]]) -> None:
    if not rows:
        return
    import csv
    with open(path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="SSL pretraining on LIVECell")
    p.add_argument("--data-root", type=Path, required=True)
    p.add_argument("--output-dir", type=Path, required=True)
    p.add_argument("--method", choices=["simclr", "supcon", "msclacl"], required=True)
    p.add_argument("--morphology-dir", type=Path, default=None,
                   help="Path to morphology_feasibility output (needed for msclacl)")
    p.add_argument("--epochs", type=int, default=100)
    p.add_argument("--batch-size", type=int, default=64)
    p.add_argument("--lr", type=float, default=0.001)
    p.add_argument("--weight-decay", type=float, default=1e-6)
    p.add_argument("--temperature", type=float, default=0.07)
    p.add_argument("--feature-dim", type=int, default=128)
    p.add_argument("--image-size", type=int, default=224)
    p.add_argument("--num-workers", type=int, default=4)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"])
    return p.parse_args()


def main() -> None:
    args = parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    args.output_dir.mkdir(parents=True, exist_ok=True)

    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)

    train_ann = discover_train_ann(args.data_root)
    print(f"Train annotation: {train_ann}")

    transform = SimCLRTransform(size=args.image_size)
    need_labels = args.method == "supcon"
    need_indices = args.method == "msclacl"

    # For MS-CLACL, we need image_id to look up similarity matrix rows
    if need_indices:

        class IndexedLiveCellImageDataset(LiveCellImageDataset):
            def __getitem__(self, index):
                with Image.open(str(self.image_paths[index])) as img:
                    img = img.convert("RGB")
                v1, v2 = self.transform(img)
                return v1, v2, self.image_ids[index]

        dataset = IndexedLiveCellImageDataset(args.data_root, train_ann, transform)
    else:
        dataset = LiveCellImageDataset(args.data_root, train_ann, transform, cell_line_labels=need_labels)

    print(f"Dataset size: {len(dataset)} images")

    # SupCon requires valid cell-line labels for every image.
    # If infer_cell_line() returned -1, those samples pollute SupCon training
    # by forming a spurious "unknown" positive cluster.
    if args.method == "supcon":
        unknown = [i for i, cl in enumerate(dataset.cell_lines) if cl < 0]
        if unknown:
            examples = [str(dataset.image_paths[i].name) for i in unknown[:5]]
            raise RuntimeError(
                f"SupCon: {len(unknown)}/{len(dataset)} images have unknown cell-line label. "
                f"Examples: {examples}. Fix infer_cell_line() before training."
            )
        print(f"SupCon label check passed: all {len(dataset)} images have valid cell-line labels")
        # Report label distribution
        from collections import Counter
        dist = Counter(dataset.cell_lines)
        print("Cell-line distribution: "
              + ", ".join(f"{CELL_LINE_NAMES[k]}={v}" for k, v in sorted(dist.items())))

    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        drop_last=True,
    )

    # Load morphology similarity matrix for MS-CLACL
    similarity_matrix = None
    id_to_row: Optional[Dict[int, int]] = None
    if args.method == "msclacl":
        if args.morphology_dir is None:
            raise ValueError("--morphology-dir required for msclacl")
        sim_path = args.morphology_dir / "image_similarity_matrix.npy"
        if not sim_path.exists():
            raise FileNotFoundError(f"Similarity matrix not found: {sim_path}")
        sim_mat = np.load(str(sim_path))
        similarity_matrix = torch.from_numpy(sim_mat.astype(np.float32))
        print(f"Loaded similarity matrix: {similarity_matrix.shape}")

        # Build image_id → matrix_row mapping from morphology features CSV
        import csv as csv_mod
        feat_csv = args.morphology_dir / "image_morphology_features.csv"
        if not feat_csv.exists():
            raise FileNotFoundError(f"Morphology features CSV not found: {feat_csv}")
        id_to_row = {}
        with open(feat_csv, "r") as f:
            reader = csv_mod.DictReader(f)
            for row_idx, row in enumerate(reader):
                img_id = int(row["image_id"])
                id_to_row[img_id] = row_idx
        print(f"Built id_to_row mapping: {len(id_to_row)} entries")

        # Verify dataset image_ids are all in the mapping
        missing = [img_id for img_id in dataset.image_ids if img_id not in id_to_row]
        if missing:
            raise RuntimeError(f"{len(missing)} dataset image_ids not found in morphology matrix (e.g. {missing[:3]})")

        # Sanity check: weight distribution on the first batch.
        # If positive_ratio is ~0 the threshold is too high (no positives).
        # If positive_ratio is ~1 the threshold is too low (loss degenerates).
        check_n = min(args.batch_size, len(dataset.image_ids))
        check_ids = dataset.image_ids[:check_n]
        check_rows = [id_to_row[i] for i in check_ids]
        check_submat = similarity_matrix[check_rows][:, check_rows]
        check_weights = torch.where(
            check_submat >= 0.5, check_submat, torch.zeros_like(check_submat)
        )
        check_weights.fill_diagonal_(0.0)
        pos_count = (check_weights > 0).sum().item()
        zero_count = (check_weights == 0).sum().item()
        total = check_weights.numel()
        pos_ratio = pos_count / total if total > 0 else 0.0
        w_mean = check_weights.mean().item()
        w_min = check_weights[check_weights > 0].min().item() if pos_count > 0 else 0.0
        w_max = check_weights.max().item()
        print(
            f"MSCLACL_WEIGHT_CHECK: positive_count={pos_count} zero_count={zero_count} "
            f"positive_ratio={pos_ratio:.4f} weight_mean={w_mean:.4f} "
            f"weight_min(pos)={w_min:.4f} weight_max={w_max:.4f}"
        )
        if pos_ratio < 0.05:
            print("WARNING: positive_ratio < 0.05, threshold may be too high — few positive pairs")
        if pos_ratio > 0.95:
            print("WARNING: positive_ratio > 0.95, threshold may be too low — loss may degenerate")

    model = SSLModel(feature_dim=args.feature_dim)
    model.to(device)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.lr,
        weight_decay=args.weight_decay,
    )

    config = {
        "method": args.method,
        "epochs": args.epochs,
        "batch_size": args.batch_size,
        "lr": args.lr,
        "weight_decay": args.weight_decay,
        "temperature": args.temperature,
        "feature_dim": args.feature_dim,
        "image_size": args.image_size,
        "dataset_size": len(dataset),
        "seed": args.seed,
        "morphology_dir": str(args.morphology_dir) if args.morphology_dir else None,
        "msclacl_soft_threshold": 0.5 if args.method == "msclacl" else None,
        "device": str(device),
    }
    (args.output_dir / "ssl_config.json").write_text(json.dumps(config, indent=2))

    print(f"Starting {args.method} pretraining, {args.epochs} epochs, batch_size={args.batch_size}")
    t0 = time.time()

    loss_rows = train_ssl(
        model, loader, optimizer, device,
        epochs=args.epochs,
        method=args.method,
        temperature=args.temperature,
        similarity_matrix=similarity_matrix,
        id_to_row=id_to_row,
    )

    elapsed = time.time() - t0
    print(f"Training finished in {elapsed:.1f}s")

    write_csv(args.output_dir / "ssl_loss.csv", loss_rows)

    # Save backbone for Mask R-CNN transfer
    backbone_path = args.output_dir / "backbone_resnet50_ssl.pth"
    torch.save(model.backbone.state_dict(), backbone_path)
    print(f"Saved backbone: {backbone_path}")

    # Save full checkpoint
    ckpt_path = args.output_dir / "ssl_checkpoint.pth"
    torch.save({
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "config": config,
    }, ckpt_path)
    print(f"Saved checkpoint: {ckpt_path}")


if __name__ == "__main__":
    main()
