"""Offline model loading and non-overlapping tiling with boundary completion."""
from __future__ import annotations

import contextlib
import sys
from pathlib import Path

from PIL import Image
import torch
import yaml

from model.LTF_v3_3 import LTFNetV3_3
from utils.tile_dataset_srm import clip_transform, crop_image_to_tiles

ROOT = Path(__file__).absolute().parents[1]
BRANCHES = ('level1', 'level2', 'fusion')


def load_config(path=None):
    path = Path(path) if path else ROOT / 'configs/eval.yaml'
    with path.expanduser().resolve().open(encoding='utf-8') as stream:
        config = yaml.safe_load(stream)
    if not isinstance(config, dict):
        raise ValueError('Configuration must be a YAML mapping')
    return config


def package_path(path):
    path = Path(path).expanduser()
    return path.resolve() if path.is_absolute() else (ROOT / path).resolve()


def model_kwargs(config):
    return dict(
        num_heads=config.get('transformer_heads', 2), num_classes=2,
        num_layers=config.get('transformer_layers', 0), dropout=config.get('dropout', .1),
        tile_heads=config.get('tile_heads', 2), tile_layers=config.get('tile_layers', 1),
        use_global_image=config.get('use_global_image', True),
        use_patch_transformer=config.get('use_patch_transformer', False),
        linear_classifier=config.get('linear_classifier', True),
        mil_heads=config.get('mil_heads', 8), fusion_hidden=config.get('fusion_hidden', 512),
        l1_cross_shuffle_p=float(config.get('v2_l1_cross_shuffle_p', .9)),
        l1_replace_ratio_range=tuple(config.get('v2_l1_replace_ratio_range', [.5, .9])),
        l1_with_replacement_p=float(config.get('v2_l1_with_replacement_p', .5)),
        l2_cross_shuffle_p=float(config.get('v2_l2_cross_shuffle_p', .9)),
        l2_replace_ratio_range=tuple(config.get('v2_l2_replace_ratio_range', [.1, .8])),
        l2_with_replacement_p=float(config.get('v2_l2_with_replacement_p', .5)),
        bag_token_shuffle_p=float(config.get('v2_bag_token_shuffle_p', .5)),
        cross_shuffle_min_keep_original=int(config.get('v2_cross_shuffle_min_keep_original', 1)),
        use_clip_intermediate_for_mil=config.get('v2_use_clip_intermediate_for_mil', False),
        clip_mil_stop_after_block_1based=int(config.get('v2_clip_mil_stop_after_block_1based', 12)),
        pool_gate_hidden=int(config.get('v3_3_pool_gate_hidden', 16)),
        use_pool_gate_fusion=config.get('v3_3_use_pool_gate_fusion', True),
        use_mil_cross_shuffle=config.get('mil_use_cross_shuffle', True),
    )


def read_core_weights(path):
    weights = torch.load(path, map_location='cpu', weights_only=True)
    if not isinstance(weights, dict) or not weights or not all(torch.is_tensor(v) for v in weights.values()):
        raise ValueError('Expected an LTF v3.3 core state_dict, such as epoch_20.pth')
    return weights


def create_model(config, weights=None, device='cpu'):
    import clip
    backbone = package_path(config['backbone'])
    if not backbone.is_file():
        raise FileNotFoundError(f'Bundled CLIP backbone not found: {backbone}')
    # Local checkpoint loading is explicit: this path never downloads a model.
    with contextlib.redirect_stdout(sys.stderr):
        clip_model, _ = clip.load(str(backbone), device='cpu', jit=False)
        model = LTFNetV3_3(clip_model.visual, **model_kwargs(config))
    del clip_model
    if weights is not None:
        model.core.load_state_dict(read_core_weights(package_path(weights)), strict=True)
    return model.to(device)


def prepare_image(path, tile_size=224):
    """Original RGB and CLIP normalization; non-overlapping blocks with boundary completion."""
    path = Path(path).expanduser().resolve()
    with Image.open(path) as source:
        image = source.convert('RGB')
    transform = clip_transform(tile_size, blur_radius=None)
    local_tiles = crop_image_to_tiles(image, tile_size)
    return torch.stack([transform(image)] + [transform(tile) for tile in local_tiles])


class LTFDetector:
    def __init__(self, config=None, weights=None, device='auto', chunk_size=256):
        self.config = dict(config) if isinstance(config, dict) else load_config(config)
        if device == 'auto':
            device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        self.device = torch.device(device)
        if self.device.type == 'cuda' and not torch.cuda.is_available():
            raise RuntimeError('CUDA is unavailable; use device="cpu" for CPU inference')
        if int(chunk_size) < 1:
            raise ValueError('chunk_size must be positive')
        torch.manual_seed(int(self.config.get('seed', 42)))
        self.weights = package_path(weights or self.config.get('weights', 'weights/epoch_20.pth'))
        self.model = create_model(self.config, self.weights, self.device).eval()
        self.model.eval_chunk_size = int(chunk_size)

    @torch.inference_mode()
    def predict(self, path, branch='fusion', threshold=.5):
        if branch not in BRANCHES:
            raise ValueError(f'branch must be one of {BRANCHES}')
        if not 0 <= threshold <= 1:
            raise ValueError('threshold must be between zero and one')
        path = Path(path).expanduser().resolve()
        images = prepare_image(path, int(self.model.input_resolution))
        self.model.eval()
        output = self.model(images.unsqueeze(0).to(self.device))
        probabilities = {}
        for name in BRANCHES:
            probability = output[f'logits_{name}'].float().softmax(-1)[0].cpu()
            if not torch.isfinite(probability).all():
                raise ValueError(f'Nonfinite {name} prediction for {path}')
            probabilities[name] = {'real': float(probability[0]), 'fake': float(probability[1])}
        score = probabilities[branch]['fake']
        label = int(score >= threshold)
        return dict(path=str(path), prediction='fake' if label else 'real', label=label,
                    branch=branch, threshold=float(threshold), score_fake=score,
                    tiling_policy="nonoverlap_then_boundary",
                    views=int(images.shape[0]), probabilities=probabilities)