"""Random-initialized author architecture, embedding denoising and COSMO heads."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
from scratch_data import *
from torch import nn
import yaml,math
from threedscriptors.configuration.architecture_config import EncoderDecoderArchitectureConfig
from threedscriptors.data_handling.sample import PreprocessedSample
from threedscriptors.training.noise_scheduler import NoiseModule,ConstantSchedule
from threedscriptors.training.pretraining import atom_denoising_loss
from train_cosmo import loss_terms,batches

CONFIG=dict(epochs=10,probes_per_epoch=5,seed=42026,architecture_source_commit='20df6fe9c61154e351d677cb0f77fd716d00beae',
    initialization='Random learned Rem3Di weights; no published or PCQM/COSMO checkpoint loaded',
    learning_rate_max=1e-4,weight_decay=1e-3,learning_rate_uncertainty=1e-3,gradient_clip=1.,
    noise_sigma=.3,denoising='Exact author atomic embedding denoising loss and gradient paths, on PCQM',
    uncertainty='sum(L_t/(2*exp(s_t))+s_t/2); raw denoising loss; initial log variance ln(1088)-0.5 for denoising, -0.5 for COSMO',
    uncertainty_initialization='Equivalent initial precision to author-denoising per channel; no modification of raw denoising function',
    task_order=['rem3di_denoising','cosmo'],cosmo_head='640 -> 256 SiLU -> 155',
    pcqm_batch_size=128,pcqm_atom_cap=5250,cosmo_batch_size=96,cosmo_atom_cap=4000,
    epoch_definition='Maximum of complete PCQM and COSMO loader lengths; shorter loader cycles; all retained training molecules seen each epoch',
    scheduler='OneCycleLR max_lr=1e-4, total optimizer steps computed from ten exact deterministic epoch plans',
    decorrelation=False,combisolv=False,mace='Frozen MACE-POLAR-1-M; raw 4096-channel float32 cache',
    normalization='Fresh train-only population invariant mean/std across PCQM and COSMO atoms',
    probe='Frozen pooled 640; solute640+solvent640+Temperature_K; exact BigSolDB rows; full75731-row Kumo Large context, 4 estimators, seed42',
    probe_schedule='ceil(k * epoch_steps / 5), k=1..5, each epoch; fifty test probes; no test-based selection',
    checkpoint_policy='All fifty probe points include all weights, optimizer, scheduler and RNG; verified Hugging Face backup before probing')


def uncertainty(losses,log_variance):
    return (losses/(2*log_variance.exp())+log_variance/2).sum()


class ScratchJoint(nn.Module):
    def __init__(self,checkpoint=None,device='cuda'):
        super().__init__()
        checkpoint=Path(checkpoint) if checkpoint is not None else None
        architecture_path=checkpoint/'post_training_architecture_config.yaml' if checkpoint is not None else OUT/'architecture.yaml'
        architecture=yaml.safe_load(architecture_path.read_text())
        architecture['mace_config']=None
        cfg=EncoderDecoderArchitectureConfig.model_validate(architecture)
        # Constructors initialize every learned tensor. Only data statistics are supplied.
        stats_path=checkpoint/'normalization.pt' if checkpoint is not None else DATA/'normalization.pt'
        stats=torch.load(stats_path,map_location='cpu',weights_only=True)
        self.pre=cfg._build_preprocessor()
        self.pre.atomic_preprocessor.double();self.pre.geometric_preprocessor.float()
        self.pre.atomic_preprocessor.invariant_normalization.set_stats(stats['mean'],stats['std'],dtype=torch.float64)
        self.encoder=cfg._build_encoder().float()
        self.denoiser=cfg.decoder_config.build().float()
        self.cosmo_head=nn.Sequential(nn.Linear(640,256),nn.SiLU(),nn.Linear(256,155))
        self.log_variance=nn.Parameter(torch.tensor([math.log(1088)-.5,-.5],dtype=torch.float32))
        self.noise=NoiseModule(ConstantSchedule(.3))
        if checkpoint is not None:self.load_parts(Path(checkpoint))
        self.to(device)

    @property
    def parts(self):
        return {'encoder.pth':self.encoder,'atomic_preprocessor.pth':self.pre.atomic_preprocessor,
            'geometric_preprocessor.pth':self.pre.geometric_preprocessor,'denoising_decoder.pth':self.denoiser,
            'cosmo_head.pth':self.cosmo_head}

    def load_parts(self,path):
        manifest=json.loads((path/'manifest.json').read_text())
        for name,module in self.parts.items():
            assert sha(path/name)==manifest['files'][name]['sha256'],name
            module.load_state_dict(torch.load(path/name,map_location='cpu',weights_only=True),strict=True)
        assert sha(path/'task_weighting.pth')==manifest['files']['task_weighting.pth']['sha256']
        state=torch.load(path/'task_weighting.pth',map_location='cpu',weights_only=True)
        with torch.no_grad():self.log_variance.copy_(state['log_variance'])

    def encode(self,sample):
        clean=self.pre(sample);descriptor=self.encoder(clean)
        z=descriptor if torch.is_tensor(descriptor) else descriptor.flat
        return clean,descriptor,z

    def denoising(self,sample):
        clean,descriptor,z=self.encode(sample)
        target=clean.preprocessed_atomic_embeddings.clone()
        noisy=self.noise(target,clean.padding_mask)
        corrupted=PreprocessedSample(preprocessed_atomic_embeddings=noisy,
            initial_pair_representation=clean.initial_pair_representation.detach(),
            geometrical_encoding=clean.geometrical_encoding.detach(),
            padding_mask=clean.padding_mask,pair_mask=clean.pair_mask)
        reconstruction=self.denoiser(corrupted,descriptor)
        loss=atom_denoising_loss(target,reconstruction,clean.padding_mask,.3)
        return loss,z,target

    def cosmo(self,sample,target):
        _,_,z=self.encode(sample);pred=self.cosmo_head(z)
        loss,parts=loss_terms(pred,target)
        return loss,z,pred,parts


def optimizer_for(model):
    return torch.optim.AdamW([
        dict(params=list(model.pre.parameters())+list(model.encoder.parameters())+list(model.denoiser.parameters())+list(model.cosmo_head.parameters()),lr=CONFIG['learning_rate_max']),
        dict(params=[model.log_variance],lr=CONFIG['learning_rate_uncertainty'],weight_decay=0.)],weight_decay=CONFIG['weight_decay'])