# These imports are tricky because they use c++, do not move them
import os, shutil
import warnings

import torch
import hydra
from omegaconf import DictConfig
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint

import utils
from datasets import dataset

from diffusion_model import GraphDiffusionTransformer
from metrics.molecular_metrics_sampling import SamplingMolecularMetrics

from analysis.visualization import MolecularVisualization

warnings.filterwarnings("ignore", category=UserWarning)
torch.set_float32_matmul_precision("medium")

def get_resume(cfg, model_kwargs):
    """Resumes a run. It loads previous config without allowing to update keys (used for testing)."""
    saved_cfg = cfg.copy()
    name = cfg.general.name + "_resume"
    resume = cfg.general.test_only
    batch_size = cfg.train.batch_size
    model = GraphDiffusionTransformer.load_from_checkpoint(resume, **model_kwargs)
    cfg = model.cfg
    cfg.general.test_only = resume
    cfg.general.name = name
    cfg.train.batch_size = batch_size
    cfg = utils.update_config_with_new_keys(cfg, saved_cfg)
    return cfg, model

def get_resume_adaptive(cfg, model_kwargs, test_only=False):
    """ Resumes a run. It loads previous config but allows to make some changes (used for resuming training)."""
    saved_cfg = cfg.copy()
    # Fetch path to this file to get base path
    current_path = os.path.dirname(os.path.realpath(__file__))
    root_dir = current_path.split('outputs')[0]

    if test_only:
        resume_path = os.path.join(root_dir, cfg.general.test_only)
    else:
        resume_path = os.path.join(root_dir, cfg.general.resume)

    # Load the checkpoint manually first to get the model weights only
    checkpoint = torch.load(resume_path)
    
    model = GraphDiffusionTransformer(cfg=cfg, **model_kwargs)
    
    # Load only the model weights from the checkpoint, not the optimizer state

    missing, unexpected = model.load_state_dict(checkpoint['state_dict'], strict=False)
    print("Missing keys:", missing)
    print("Unexpected keys:", unexpected)
    print('model loaded')
    
    return saved_cfg, model

@hydra.main(
    version_base="1.1", config_path="../configs", config_name="config"
)
def main(cfg: DictConfig):
    datamodule = dataset.DataModule(cfg)
    datamodule.prepare_split()
    tokenizer = datamodule.prepare_tokenizer(cfg)
    datamodule.prepare_data()
    dataset_infos = dataset.DataInfos(datamodule=datamodule, cfg=cfg, tokenizer=tokenizer)
    # train_smiles, reference_smiles, sampling_condition_dict = datamodule.get_train_smiles()
    vocab_smiles = datamodule.get_node_vocabulary(cfg)

    dataset_infos.compute_input_output_dims(datamodule=datamodule)
    # train_metrics = TrainMolecularMetricsDiscrete(dataset_infos)

    sampling_metrics = SamplingMolecularMetrics(dataset_infos, vocab_smiles)
    visualization_tools = MolecularVisualization(dataset_infos)

    model_kwargs = {
        "dataset_infos": dataset_infos,
        "sampling_metrics": sampling_metrics,
        "visualization_tools": visualization_tools,
        "tokenizer": tokenizer,
    }

    if cfg.general.test_only:
        # When testing, previous configuration is fully loaded
        # cfg, _ = get_resume(cfg, model_kwargs)
        cfg, model = get_resume_adaptive(cfg, model_kwargs, test_only=True)
        os.chdir(cfg.general.test_only.split("checkpoints")[0])
    elif cfg.general.resume is not None:
        # When resuming, we can override some parts of previous configuration
        cfg, model = get_resume_adaptive(cfg, model_kwargs)
        os.chdir(cfg.general.resume.split("checkpoints")[0])
    else:
        model = GraphDiffusionTransformer(cfg=cfg, **model_kwargs)

    # Define checkpoint callback
    callbacks = []
    if cfg.general.save_model:
        checkpoint_callback = ModelCheckpoint(
            dirpath=f"checkpoints/{cfg.general.name}",
            filename="last",
            save_last=True,
            save_top_k=1,
            enable_version_counter=False,
            verbose=True,
        )
        callbacks.append(checkpoint_callback)

    trainer = Trainer(
        gradient_clip_val=cfg.train.clip_grad,
        accelerator="gpu"
        if torch.cuda.is_available() and cfg.general.gpus > 0
        else "cpu",
        devices=cfg.general.gpus
        if torch.cuda.is_available() and cfg.general.gpus > 0
        else None,
        max_epochs=cfg.train.n_epochs,
        # fast_dev_run=5,
        # limit_train_batches=5,
        # limit_val_batches=5,
        # limit_test_batches=2,
        # precision=16,
        enable_checkpointing=cfg.general.save_model,
        check_val_every_n_epoch=cfg.train.check_val_every_n_epoch,
        val_check_interval=cfg.train.val_check_interval,
        strategy="ddp" if cfg.general.gpus > 1 else "auto",
        enable_progress_bar=cfg.general.enable_progress_bar,
        callbacks=callbacks,
        reload_dataloaders_every_n_epochs=0,
        logger=[],
    )

    if not cfg.general.test_only:
        # trainer.fit(model, datamodule=datamodule, ckpt_path=cfg.general.resume)
        trainer.fit(model, datamodule=datamodule)
        if cfg.general.save_model:
            trainer.save_checkpoint(f"checkpoints/{cfg.general.name}/last.ckpt")
        trainer.test(model, datamodule=datamodule)
    else:
        trainer.test(model, datamodule=datamodule, ckpt_path=cfg.general.test_only)

if __name__ == "__main__":
    main()
