"""Load the study's native SCimilarity encoder and decoder with PyTorch."""

import json
from pathlib import Path

import torch

from _scimilarity_networks import Decoder, Encoder


class SCimilarity(torch.nn.Module):
    def __init__(self, gene_order, hidden_dim, latent_dim):
        super().__init__()
        self.gene_order = tuple(gene_order)
        self.n_genes = len(self.gene_order)
        self.latent_dim = latent_dim
        self.encoder = Encoder(self.n_genes, latent_dim=latent_dim, hidden_dim=hidden_dim)
        self.decoder = Decoder(self.n_genes, latent_dim=latent_dim, hidden_dim=hidden_dim)

    def forward(self, expression):
        embedding = self.encoder(expression)
        return embedding, self.decoder(embedding)


def _load_state(path):
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    if not isinstance(checkpoint, dict) or set(checkpoint) != {"state_dict"}:
        raise ValueError(f"Expected only a state_dict in {path}")
    state = checkpoint["state_dict"]
    if not isinstance(state, dict) or not state:
        raise ValueError(f"Missing model tensors in {path}")
    if any(not isinstance(value, torch.Tensor) or not torch.isfinite(value).all()
           for value in state.values()):
        raise ValueError(f"Non-tensor or nonfinite model state in {path}")
    return state


def load_scimilarity(model_directory):
    directory = Path(model_directory)
    gene_order = [
        line.strip() for line in (directory / "gene_order.tsv").read_text().splitlines()
    ]
    if not gene_order or any(not gene for gene in gene_order) or len(set(gene_order)) != len(gene_order):
        raise ValueError("SCimilarity gene order must contain unique, nonempty gene symbols")
    encoder_state = _load_state(directory / "encoder.ckpt")
    layer_sizes = json.loads((directory / "layer_sizes.json").read_text())
    if {name: list(value.shape) for name, value in encoder_state.items()} != layer_sizes:
        raise ValueError("SCimilarity encoder tensors disagree with layer_sizes.json")
    layers = [
        shape for name, shape in sorted(
            layer_sizes.items(), key=lambda item: int(item[0].split(".")[1])
        )
        if name.endswith(".weight") and len(shape) == 2
    ]
    if len(layers) < 2 or layers[0][1] != len(gene_order):
        raise ValueError("SCimilarity network dimensions disagree with the gene order")
    model = SCimilarity(
        gene_order, hidden_dim=[shape[0] for shape in layers[:-1]],
        latent_dim=layers[-1][0],
    )
    model.encoder.load_state_dict(encoder_state, strict=True)
    model.decoder.load_state_dict(_load_state(directory / "decoder.ckpt"), strict=True)
    model.eval()
    return model