"""Research recognizer sharing one English encoder across two explicit branches.

The character branch applies saved low-rank corrections through functional_call
and supplies precomputed encoder outputs to a decoder-only VED container. No
second learned encoder is registered or saved. Functional parameter overrides
are serial CPU research operations, not a concurrent/native serving interface.
"""
import copy
import hashlib
import json
from pathlib import Path

import torch
from torch import nn
from safetensors.torch import load_file
from transformers import DeiTModel, GenerationConfig, VisionEncoderDecoderConfig, VisionEncoderDecoderModel


def sha(path):
    value = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''): value.update(block)
    return value.hexdigest()


def eager(config):
    config._attn_implementation = 'eager'
    config.encoder._attn_implementation = 'eager'
    config.decoder._attn_implementation = 'eager'
    return config


def load_english(folder):
    config = eager(VisionEncoderDecoderConfig.from_pretrained(folder, local_files_only=True))
    model, report = VisionEncoderDecoderModel.from_pretrained(folder, config=config,
        encoder=DeiTModel(config.encoder, add_pooling_layer=False), local_files_only=True,
        dtype=torch.float32, attn_implementation='eager', output_loading_info=True)
    assert not any(report.get(k) for k in ['missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs'])
    assert sum(p.numel() for p in model.parameters()) == 61448832
    return model.eval().requires_grad_(False)


class PrecomputedOnlyEncoder(nn.Module):
    main_input_name = 'pixel_values'

    def __init__(self, config):
        super().__init__(); self.config = copy.deepcopy(config)

    def get_output_embeddings(self):
        return None

    def forward(self, pixel_values=None, **kwargs):
        raise RuntimeError('Character decoding requires the shared encoder outputs')


class Correction(nn.Module):
    def __init__(self, spec):
        super().__init__(); self.spec = spec
        if spec['kind'] == 'factor':
            self.left = nn.Parameter(torch.empty(spec['matrix_shape'][0], spec['rank']))
            self.right = nn.Parameter(torch.empty(spec['rank'], spec['matrix_shape'][1]))
        else:
            assert spec['kind'] == 'replacement'
            self.value = nn.Parameter(torch.empty(spec['shape']))

    def apply(self, base):
        if self.spec['kind'] == 'replacement': return self.value
        return base + (self.left @ self.right).reshape(self.spec['shape'])


class SharedCharacterRecognizer(nn.Module):
    def __init__(self, original, character_decoder, character_projection, character_config,
                 english_generation, character_generation, specifications):
        super().__init__()
        self.original = original
        cfg = eager(copy.deepcopy(character_config))
        self.character = VisionEncoderDecoderModel(config=cfg,
            encoder=PrecomputedOnlyEncoder(cfg.encoder), decoder=character_decoder)
        if character_projection is not None: self.character.enc_to_dec_proj = character_projection
        self.original.generation_config = copy.deepcopy(english_generation)
        self.character.generation_config = copy.deepcopy(character_generation)
        self.specifications = specifications
        self.corrections = nn.ModuleDict({spec['key']: Correction(spec) for spec in specifications})

    def generate_pixels(self, pixels, branch, generation):
        if self.training: raise ValueError('Explicit eval mode is required')
        if branch == 'english': return self.original.generate(pixel_values=pixels, **generation)
        if branch != 'character': raise ValueError('Require an explicit English or character branch')
        base = self.original.encoder.state_dict()
        overrides = {spec['name']: self.corrections[spec['key']].apply(base[spec['name']])
                     for spec in self.specifications}
        output = torch.func.functional_call(self.original.encoder, overrides, (),
            dict(pixel_values=pixels, return_dict=True), strict=False)
        return self.character.generate(encoder_outputs=output, **generation)

    @classmethod
    def from_bundle(cls, folder):
        folder = Path(folder); manifest = json.loads((folder / 'bundle.json').read_text())
        for name, expected in manifest['files_sha256'].items(): assert sha(folder / name) == expected
        configs = json.loads((folder / 'model-configs.json').read_text())
        english = eager(VisionEncoderDecoderConfig.from_dict(configs['english']))
        character = eager(VisionEncoderDecoderConfig.from_dict(configs['character']))
        original = VisionEncoderDecoderModel(config=english,
            encoder=DeiTModel(english.encoder, add_pooling_layer=False))
        private = VisionEncoderDecoderModel(config=character,
            encoder=PrecomputedOnlyEncoder(character.encoder))
        model = cls(original, private.decoder, getattr(private, 'enc_to_dec_proj', None), character,
            GenerationConfig.from_pretrained(folder, config_file_name='english-generation.json', local_files_only=True),
            GenerationConfig.from_pretrained(folder, config_file_name='character-generation.json', local_files_only=True),
            manifest['corrections'])
        result = model.load_state_dict(load_file(str(folder / 'weights.safetensors')), strict=True)
        assert not result.missing_keys and not result.unexpected_keys
        assert sum(p.numel() for p in model.parameters()) == manifest['parameters'] == 74842752
        assert len(model.character.encoder.state_dict()) == 0
        return model.eval().requires_grad_(False)