# Copyright Generate Biomedicines, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Utilities to save and load models with metadata.
"""

import os
import os.path as osp
import tempfile
from pathlib import Path
from urllib.parse import parse_qs, urlparse
from uuid import uuid4

import torch

import chroma.utility.api as api
from chroma.constants.named_models import NAMED_MODELS


def save_model(model, weight_file, metadata=None):
    """Save model, including optional metadata.

    Args:
        model (nn.Module): The model to save. Details about the model needed
            for initialization, such as layer sizes, should be in model.kwargs.
        weight_file (str): The destination path for saving model weights.
        metadata (dict): A dictionary of additional metadata to add to the model
            weights. For example, when saving models during training it can be
            useful to store `args` representing the CLI args, the date and time
            of training, etc.
    """
    save_dict = {"init_kwargs": model.kwargs, "model_state_dict": model.state_dict()}
    if metadata is not None:
        save_dict.update(metadata)
    local_path = str(
        Path(tempfile.gettempdir(), str(uuid4())[:8])
        if weight_file.startswith("s3:")
        else weight_file
    )
    torch.save(save_dict, local_path)
    if weight_file.startswith("s3:"):
        raise NotImplementedError("Uploading to an s3 link not supported.")


def load_model(
    weights,
    model_class,
    device="cpu",
    strict=False,
    strict_unexpected=True,
    verbose=True,
):
    """Load model saved with save_model.

    Args:
        weights (str): The destination path of the model weights to load.
            Compatible with files saved by `save_model`.
        model_class: Name of model class.
        device (str, optional): Pytorch device specification, e.g. `'cuda'` for
        GPU. Default is `'cpu'`.
        strict (bool): Whether to require that the keys match between the
            input file weights and the model created from the parameters stored
            in the model kwargs.
        strict_unexpected (bool): Whether to require that there are no
            unexpected keys when loading model weights, as distinct from the
            strict option which doesn't allow for missing keys either. By
            default, we use this option rather than strict for ease of
            development when adding model features.
        verbose (bool, optional): Show outputs from download and loading. Default True.

    Returns:
        model (nn.Module): Torch model with loaded weights.
    """

    # Process weights path
    if str(weights).startswith("named:"):
        weights = weights.split("named:")[1]
        if weights not in NAMED_MODELS[model_class.__name__]:
            raise Exception(f"Unknown {model_class.__name__} model name: {weights},")
        weights = NAMED_MODELS[model_class.__name__][weights]["s3_uri"]

    # resolve s3 paths
    if str(weights).startswith("s3:"):
        raise NotImplementedError("Loading Models from an S3 link not supported.")

    # download public models from generate
    if str(weights).startswith("https:"):
        # Decompose into arguments
        parsed_url = urlparse(weights)
        base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{parsed_url.path}"
        model_name = parse_qs(parsed_url.query).get("weights", [None])[0]
        weights = api.download_from_generate(
            base_url, model_name, force=False, exist_ok=True
        )

    # load model weights
    params = torch.load(weights, map_location="cpu")
    model = model_class(**params["init_kwargs"]).to(device)
    missing_keys, unexpected_keys = model.load_state_dict(
        params["model_state_dict"], strict=strict
    )
    if strict_unexpected and len(unexpected_keys) > 0:
        raise Exception(
            f"Error loading model from checkpoint file: {weights} contains {len(unexpected_keys)} unexpected keys: {unexpected_keys}"
        )
    return model