"""Load synthetic-pretrained Voronoi microstructure ResNets without pickle."""
from pathlib import Path

import torch
from torch import nn
from torchvision import models
from safetensors.torch import load_file

TARGET_NAMES = [
    'log_grains', 'mean_ecd', 'area_weighted_ecd', 'ecd_cv',
    'hall_petch_mean_ecd', 'phase_area_fraction', 'phase_grain_fraction',
    'log_stretch', 'log_defects', 'defect_area_fraction',
    'mean_defect_radius', 'boundary_fraction', 'interior_grain_fraction',
]
VARIANTS = {
    'resnet18_seed1019': 'resnet18',
    'resnet18_seed2026': 'resnet18',
    'resnet34_seed1019': 'resnet34',
    'resnet50_seed1019': 'resnet50',
}


class VoronoiResNet(nn.Module):
    def __init__(self, architecture):
        super().__init__()
        if architecture not in {'resnet18', 'resnet34', 'resnet50'}:
            raise ValueError('Expected resnet18, resnet34, or resnet50')
        self.encoder = getattr(models, architecture)(weights=None)
        dimension = self.encoder.fc.in_features
        self.encoder.fc = nn.Identity()
        self.head = nn.Sequential(nn.Linear(dimension, 128), nn.SiLU(), nn.Linear(128, 26))
        self.register_buffer('mean', torch.tensor([.485, .456, .406])[None, :, None, None])
        self.register_buffer('scale', torch.tensor([.229, .224, .225])[None, :, None, None])
        self.register_buffer('target_mean', torch.zeros(13))
        self.register_buffer('target_scale', torch.ones(13))

    def forward(self, images):
        """Return normalized means, clamped log variances, and encoder embeddings."""
        if images.ndim != 4 or images.shape[1:] != (1, 256, 256):
            raise ValueError('Expected floating point images shaped [N, 1, 256, 256]')
        if not images.is_floating_point():
            raise ValueError('Convert image pixels to floating point in [0, 1] first')
        z = self.encoder((images.repeat(1, 3, 1, 1) - self.mean) / self.scale)
        p = self.head(z)
        return p[:, :13], p[:, 13:].clamp(-4, 3), z

    @torch.inference_mode()
    def predict(self, images, tta=True):
        """Return parameter estimates, standard deviations, and mean embeddings.

        Log-transformed targets stay in their named log domain. Uncertainty is
        estimated by the synthetic model and is not calibrated on real images.
        """
        if self.training:
            raise ValueError('Call model.eval() before inference')
        means, variances, embeddings = [], [], []
        views = ((flip, rotation) for flip in (False, True) for rotation in range(4)) if tta else [(False, 0)]
        for flip, rotation in views:
            x = images.flip(-1) if flip else images
            p, lv, z = self(torch.rot90(x, rotation, (-2, -1)))
            means.append(p.float())
            variances.append(lv.float().exp())
            embeddings.append(z.float())
        p = torch.stack(means)
        parameters = p.mean(0) * self.target_scale + self.target_mean
        variance = torch.stack(variances).mean(0) + p.var(0, unbiased=False)
        uncertainty = variance.sqrt() * self.target_scale
        return parameters, uncertainty, torch.stack(embeddings).mean(0)


def load_model(directory, variant='resnet18_seed1019', device='cpu'):
    """Load one variant from a downloaded repository directory."""
    model = VoronoiResNet(VARIANTS[variant])
    model.load_state_dict(load_file(str(Path(directory) / f'{variant}.safetensors')), strict=True)
    return model.to(device).eval()


if __name__ == '__main__':
    torch.set_num_threads(2)
    model = load_model(Path(__file__).resolve().parent)
    p, u, z = model.predict(torch.full((1, 1, 256, 256), .6), tta=False)
    assert p.shape == u.shape == (1, 13) and z.shape == (1, 512)
    assert torch.isfinite(p).all() and torch.isfinite(u).all() and (u > 0).all()
    print('Load and inference check passed')