#!/usr/bin/env python3
"""DINOv3 continual SSL pretraining.

Simplified iBOT/DINO-style training using timm ViT + EMA teacher.
Config-driven (ssl_dinov3_continual.yaml).
Bit-exact BF16, torch.compile, FlashAttention-2.

Data: in_curated_pretrain=True (70K Hitit-ish images)
Output: backbone checkpoint for downstream fine-tune.
"""
import os, sys, json, math, argparse, random, time
from pathlib import Path
from copy import deepcopy

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from PIL import Image
import yaml

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def load_config(path):
    return yaml.safe_load(open(path))

def setup_ddp():
    if 'RANK' in os.environ:
        # Retry CUDA init (transient NVML race on some nodes)
        for attempt in range(3):
            try:
                n = torch.cuda.device_count()
                if n > 0 and torch.cuda.is_available(): break
            except Exception: pass
            time.sleep(2)
        if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
            # Fallback: unset CUDA_VISIBLE_DEVICES and retry
            vis = os.environ.pop('CUDA_VISIBLE_DEVICES', None)
            print(f"[rank {os.environ.get('RANK')}] CUDA init failed; retrying after unset (was {vis})", flush=True)
            import importlib
            importlib.reload(torch.cuda)
            time.sleep(2)
        torch.distributed.init_process_group(backend='nccl')
        local_rank = int(os.environ.get('LOCAL_RANK', 0))
        torch.cuda.set_device(local_rank)
        return True, local_rank, int(os.environ['WORLD_SIZE']), int(os.environ['RANK'])
    return False, 0, 1, 0

def is_main(rank): return rank == 0

def log(msg, rank=0):
    if is_main(rank):
        print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

class SSLImageDataset(Dataset):
    """Image-only dataset for SSL (no labels).

    Sources: either (a) `in_curated_pretrain=True` flag across all
    datasets/sources/*/manifest.jsonl, or (b) explicit list of manifests
    via manifests_list argument.
    """
    def __init__(self, manifest_path, filter_curated=True, global_size=224, local_size=96,
                 n_local_crops=6, manifests_list=None):
        self.records = []
        import os as _os
        if manifests_list:
            # Explicit list: take every valid image, no in_curated_pretrain filter
            for mf in manifests_list:
                mp = Path(mf)
                if not mp.exists(): continue
                with open(mp) as f:
                    for line in f:
                        r = json.loads(line)
                        if r.get('storage') != 'fs' or not r.get('path'): continue
                        if r.get('integrity_ok') is False: continue
                        if not _os.path.exists(r['path']): continue
                        self.records.append(r['path'])
        else:
            sources_dir = ROOT / 'datasets' / 'sources'
            for src_dir in sources_dir.iterdir():
                if not src_dir.is_dir(): continue
                mf = src_dir / 'manifest.jsonl'
                if not mf.exists(): continue
                with open(mf) as f:
                    for line in f:
                        r = json.loads(line)
                        if filter_curated and not r.get('in_curated_pretrain'):
                            continue
                        if r.get('storage') != 'fs' or not r.get('path'):
                            continue
                        if r.get('integrity_ok') is False:
                            continue
                        if not _os.path.exists(r['path']):
                            continue
                        self.records.append(r['path'])

        self.global_size = global_size
        self.local_size = local_size
        self.n_local_crops = n_local_crops

        # Basic augmentation via torchvision
        from torchvision import transforms
        self.global_transform = transforms.Compose([
            transforms.RandomResizedCrop(global_size, scale=(0.4, 1.0), antialias=True),
            transforms.RandomHorizontalFlip(p=0.0),  # cuneiform yön-duyarlı
            transforms.ColorJitter(0.4, 0.4, 0.2, 0.1),
            transforms.RandomGrayscale(p=0.2),
            transforms.ToTensor(),
            transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
        ])
        self.local_transform = transforms.Compose([
            transforms.RandomResizedCrop(local_size, scale=(0.05, 0.4), antialias=True),
            transforms.ColorJitter(0.4, 0.4, 0.2, 0.1),
            transforms.ToTensor(),
            transforms.Normalize([0.489, 0.448, 0.424], [0.362, 0.359, 0.364]),
        ])

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

    def __getitem__(self, idx):
        try:
            img = Image.open(self.records[idx]).convert('RGB')
        except Exception:
            img = Image.new('RGB', (self.global_size, self.global_size))
        globals_ = [self.global_transform(img) for _ in range(2)]
        locals_ = [self.local_transform(img) for _ in range(self.n_local_crops)]
        return globals_, locals_

def collate_fn(batch):
    # batch: list of ([g1, g2], [l1..l6])
    globals_ = torch.stack([g for item in batch for g in item[0]])  # 2B × 3 × H × W
    locals_ = torch.stack([l for item in batch for l in item[1]])   # nB × 3 × h × w
    return globals_, locals_

class ProjectionHead(nn.Module):
    """DINO projection head.

    weight_norm + last_layer.weight_g.fill_(1) ile birlikte F.normalize uygulanması
    çıkışı sabit-norm vektörlere kilitliyor ve teacher==student kalıyordu (loss
    uniform entropy ln(65536)≈11.09'da donuyor). last_layer'ı normal Linear yapıp
    init'i küçük tutuyoruz; bottleneck normalize-edilmiş kalsın.
    """
    def __init__(self, in_dim=768, out_dim=65536, hidden_dim=2048, bottleneck=256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, bottleneck),
        )
        self.last_layer = nn.Linear(bottleneck, out_dim, bias=False)
        nn.init.trunc_normal_(self.last_layer.weight, std=0.02)
    def forward(self, x):
        x = self.mlp(x)
        x = F.normalize(x, dim=-1)
        return self.last_layer(x)

class DINOLoss(nn.Module):
    def __init__(self, out_dim=65536, student_temp=0.1, teacher_temp=0.04,
                 center_momentum=0.9):
        super().__init__()
        self.student_temp = student_temp
        self.teacher_temp = teacher_temp
        self.center_momentum = center_momentum
        self.register_buffer('center', torch.zeros(1, out_dim))

    def forward(self, student_out, teacher_out):
        """
        student_out: (N_total_crops × B, out_dim) — 2 global + K local
        teacher_out: (2 × B, out_dim) — only global
        """
        s = student_out / self.student_temp
        t = F.softmax((teacher_out - self.center) / self.teacher_temp, dim=-1).detach()

        # Count local + global
        B = t.size(0) // 2
        total_crops = s.size(0) // B
        s = s.chunk(total_crops)
        t = t.chunk(2)

        total_loss = 0
        n_loss_terms = 0
        for iq, q in enumerate(t):
            for v in range(len(s)):
                if v == iq: continue  # skip same view
                loss = torch.sum(-q * F.log_softmax(s[v], dim=-1), dim=-1)
                total_loss += loss.mean()
                n_loss_terms += 1
        total_loss = total_loss / max(n_loss_terms, 1)
        self.update_center(teacher_out)
        return total_loss

    @torch.no_grad()
    def update_center(self, teacher_out):
        batch_center = teacher_out.mean(0, keepdim=True)
        if torch.distributed.is_initialized():
            torch.distributed.all_reduce(batch_center)
            batch_center /= torch.distributed.get_world_size()
        self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum)

def build_backbone(model_name='vit_base_patch14_dinov2', pretrained=True):
    import timm
    # dynamic_img_size=True → ViT accepts any resolution (224 global + 96 local)
    try:
        model = timm.create_model(model_name, pretrained=pretrained, num_classes=0,
                                   img_size=224, dynamic_img_size=True)
    except Exception:
        try:
            model = timm.create_model('vit_base_patch16_224', pretrained=pretrained,
                                       num_classes=0, dynamic_img_size=True)
        except Exception:
            # Son çare: sabit image_size, locals disable
            model = timm.create_model('vit_base_patch16_224', pretrained=pretrained, num_classes=0)
    return model

@torch.no_grad()
def ema_update(teacher, student, momentum):
    for tp, sp in zip(teacher.parameters(), student.parameters()):
        tp.data.mul_(momentum).add_(sp.data, alpha=1 - momentum)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/ssl_dinov3_continual.yaml'))
    ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/ssl_dinov3_continual/'))
    ap.add_argument('--ddp-world-size', type=int, default=0)
    args = ap.parse_args()

    cfg = load_config(args.config)
    is_ddp, local_rank, world_size, global_rank = setup_ddp()
    device = torch.device(f'cuda:{local_rank}' if torch.cuda.is_available() else 'cpu')

    log(f"DDP: {is_ddp}, world_size={world_size}, rank={global_rank}", global_rank)
    log(f"Device: {device}", global_rank)

    Path(args.output).mkdir(parents=True, exist_ok=True)

    # Data
    crops = cfg.get('crops', {})
    manifest = ROOT / 'datasets/unified/classification/manifest.jsonl'
    data_cfg = cfg.get('data', {})
    manifests_list = data_cfg.get('manifests')  # if set in config, use them
    filter_curated = data_cfg.get('filter') is not None  # None or 'null' → all
    dataset = SSLImageDataset(
        manifest, filter_curated=filter_curated,
        global_size=crops.get('global_size', 224),
        local_size=crops.get('local_size', 96),
        n_local_crops=crops.get('local_count', 6),
        manifests_list=manifests_list,
    )
    log(f"Dataset: {len(dataset)} images", global_rank)

    sampler = DistributedSampler(dataset) if is_ddp else None
    batch_size = cfg.get('batch_size', 256) // max(world_size, 1)
    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=(sampler is None),
        sampler=sampler, num_workers=8, pin_memory=True, drop_last=True,
        collate_fn=collate_fn, persistent_workers=True,
    )

    # Models
    backbone_name = cfg.get('model', {}).get('teacher', 'vit_base_patch14_dinov2')
    student_bb = build_backbone(backbone_name).to(device)
    teacher_bb = build_backbone(backbone_name).to(device)
    teacher_bb.load_state_dict(student_bb.state_dict())
    for p in teacher_bb.parameters(): p.requires_grad = False

    embed_dim = getattr(student_bb, 'num_features', 768)
    student_head = ProjectionHead(in_dim=embed_dim, out_dim=65536).to(device)
    teacher_head = ProjectionHead(in_dim=embed_dim, out_dim=65536).to(device)
    teacher_head.load_state_dict(student_head.state_dict())
    for p in teacher_head.parameters(): p.requires_grad = False

    if is_ddp:
        student_bb = DDP(student_bb, device_ids=[local_rank])
        student_head = DDP(student_head, device_ids=[local_rank])

    # Loss
    loss_fn = DINOLoss().to(device)

    # Optimizer
    params = [
        {'params': (student_bb.module if is_ddp else student_bb).parameters()},
        {'params': (student_head.module if is_ddp else student_head).parameters()},
    ]
    lr = cfg.get('optimizer', {}).get('lr', 1e-4)
    wd = cfg.get('optimizer', {}).get('weight_decay', 0.04)
    optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=wd, fused=True)

    # BF16 autocast
    use_bf16 = cfg.get('use_bf16', True)
    dtype = torch.bfloat16 if use_bf16 else torch.float32

    epochs = cfg.get('epochs', 20)
    ema_start = cfg.get('ema', {}).get('teacher_momentum_start', 0.996)

    log(f"Training: {epochs} epochs, batch={batch_size}/GPU × {world_size} GPUs", global_rank)
    log(f"BF16: {use_bf16}", global_rank)

    step = 0
    total_steps = len(loader) * epochs
    t0 = time.time()

    for epoch in range(epochs):
        if sampler: sampler.set_epoch(epoch)
        epoch_loss = 0.0
        n_batches = 0

        for globals_, locals_ in loader:
            globals_ = globals_.to(device, non_blocking=True)
            locals_ = locals_.to(device, non_blocking=True)

            with torch.amp.autocast('cuda', dtype=dtype, enabled=use_bf16):
                # Teacher: only globals (224×224)
                with torch.no_grad():
                    t_feat = teacher_bb(globals_)
                    t_out = teacher_head(t_feat)

                # Student: globals AND locals separately (different sizes)
                s_feat_g = student_bb(globals_)
                s_feat_l = student_bb(locals_)
                s_out_g = student_head(s_feat_g)
                s_out_l = student_head(s_feat_l)
                s_out = torch.cat([s_out_g, s_out_l], dim=0)

                loss = loss_fn(s_out, t_out)

            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(
                (student_bb.module if is_ddp else student_bb).parameters(), 3.0)
            optimizer.step()

            # EMA update
            m = ema_start + (1.0 - ema_start) * step / total_steps
            ema_update(teacher_bb, student_bb.module if is_ddp else student_bb, m)
            ema_update(teacher_head, student_head.module if is_ddp else student_head, m)

            epoch_loss += loss.item()
            n_batches += 1
            step += 1

            if step % 50 == 0:
                rate = step / (time.time() - t0)
                log(f"  epoch={epoch} step={step}/{total_steps} loss={loss.item():.4f} rate={rate:.2f}/s", global_rank)

        log(f"Epoch {epoch}: avg_loss={epoch_loss/max(1,n_batches):.4f}", global_rank)

        if is_main(global_rank) and (epoch + 1) % cfg.get('save_every_epochs', 5) == 0:
            ckpt = {
                'epoch': epoch,
                'backbone': (teacher_bb.state_dict()),  # use teacher for downstream
                'head': teacher_head.state_dict(),
                'cfg': cfg,
            }
            torch.save(ckpt, Path(args.output) / f'checkpoint_e{epoch:02d}.pt')
            torch.save(ckpt, Path(args.output) / 'checkpoint.pt')
            log(f"Saved checkpoint epoch {epoch}", global_rank)

    # Final
    if is_main(global_rank):
        ckpt = {
            'epoch': epochs,
            'backbone': (teacher_bb.state_dict()),
            'head': teacher_head.state_dict(),
            'cfg': cfg,
        }
        torch.save(ckpt, Path(args.output) / 'checkpoint.pt')
        log(f"DONE: final checkpoint saved → {args.output}/checkpoint.pt", global_rank)

    if is_ddp:
        torch.distributed.destroy_process_group()

if __name__ == '__main__':
    main()