#!/usr/bin/env python3
"""Train-time augmentation pipeline (Albumentations-based).
SOTA recipe: elastic + grid distortion + morph + CLAHE + color jitter.
Reference: Albumentations OCR recipe, arXiv:2508.11499 historical HTR.
"""
from pathlib import Path
import yaml

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

def build_train_transform(config=None):
    """Returns Albumentations.Compose — import edilir training'de."""
    import albumentations as A
    from albumentations.pytorch import ToTensorV2
    
    if config is None:
        config = yaml.safe_load(open(ROOT / "hitit_ocr" / "configs" / "preprocessing.yaml"))
    aug = config.get('augment_train', {})
    norm = config.get('normalization', {})
    
    # Pipeline adımları
    steps = []
    
    # Geometric distortion (elastic + grid)
    geo_aug = []
    if aug.get('elastic_transform', {}).get('enabled'):
        et = aug['elastic_transform']
        geo_aug.append(A.ElasticTransform(
            alpha=et.get('alpha', 40.0),
            sigma=et.get('sigma', 6.0),
            alpha_affine=et.get('alpha_affine', 8.0),
            p=1.0,
            border_mode=0
        ))
    if aug.get('grid_distortion', {}).get('enabled'):
        gd = aug['grid_distortion']
        geo_aug.append(A.GridDistortion(
            num_steps=gd.get('num_steps', 5),
            distort_limit=gd.get('distort_limit', 0.15),
            p=1.0,
            border_mode=0
        ))
    if geo_aug:
        steps.append(A.OneOf(geo_aug, p=0.5))
    
    # Morphological (wedge thickness variation)
    if aug.get('morphological', {}).get('enabled'):
        m = aug['morphological']
        # Custom implementation (Albumentations'da tam yok)
        # Kullanım: training script'te manuel apply
        pass
    
    # Color jitter
    if aug.get('color_jitter', {}).get('enabled'):
        cj = aug['color_jitter']
        steps.append(A.ColorJitter(
            brightness=cj.get('brightness', 0.2),
            contrast=cj.get('contrast', 0.2),
            saturation=cj.get('saturation', 0.1),
            hue=0.0,
            p=cj.get('p', 0.5)
        ))
    
    # Random horizontal flip — cuneiform için DEĞİL (sign orientation önemli)
    # if aug.get('horizontal_flip', {}).get('enabled'):
    #     steps.append(A.HorizontalFlip(p=0.5))
    
    # Normalization
    if norm.get('strategy') == 'dataset_specific':
        # norm_stats'tan yükle
        import json
        stats_path = ROOT / "datasets" / "processed" / "normalization_stats.json"
        if stats_path.exists():
            ns = json.load(open(stats_path))
            mean = ns['global_rgb']['mean']
            std = ns['global_rgb']['std']
        else:
            mean = [0.485, 0.456, 0.406]
            std = [0.229, 0.224, 0.225]
    elif norm.get('strategy') == 'imagenet':
        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
    else:
        mean = [0.0, 0.0, 0.0]
        std = [1.0, 1.0, 1.0]
    
    steps.append(A.Normalize(mean=mean, std=std))
    steps.append(ToTensorV2())
    
    return A.Compose(steps)

def build_val_transform(config=None):
    """Validation/test: sadece normalize."""
    import albumentations as A
    from albumentations.pytorch import ToTensorV2
    import json
    
    if config is None:
        config = yaml.safe_load(open(ROOT / "hitit_ocr" / "configs" / "preprocessing.yaml"))
    norm = config.get('normalization', {})
    
    if norm.get('strategy') == 'dataset_specific':
        stats_path = ROOT / "datasets" / "processed" / "normalization_stats.json"
        if stats_path.exists():
            ns = json.load(open(stats_path))
            mean, std = ns['global_rgb']['mean'], ns['global_rgb']['std']
        else:
            mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
    else:
        mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
    
    return A.Compose([
        A.Normalize(mean=mean, std=std),
        ToTensorV2(),
    ])

if __name__ == '__main__':
    # Test/print pipeline
    try:
        import albumentations
        t = build_train_transform()
        print(f"Train pipeline: {len(t.transforms)} steps")
        for step in t.transforms:
            print(f"  - {type(step).__name__}: p={getattr(step, 'p', 'n/a')}")
    except ImportError:
        print("albumentations kurulu değil. pip install albumentations")