"""
model.py
========
Public entry point for WEO-SAS/sen2sr stored on HuggingFace Hub.

All parameters are read from config.json.

Usage
-----
    from huggingface_hub import snapshot_download
    import sys

    local_dir = snapshot_download("WEO-SAS/sen2sr")
    sys.path.insert(0, local_dir)
    from model import Model

    model = Model(local_dir=local_dir)

    # Array inference: (4, H, W) float32 in [0, 1] -> (4, H*4, W*4) float32
    sr = model.predict(image)

    # GeoTIFF pipeline
    model.predict_tif("s2_scene.tif", "s2_sr.tif")
"""

from __future__ import annotations

import importlib.util
import json
import os
import sys
from typing import List, Optional

import numpy as np


def _load_module(name: str, path: str):
    spec   = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


class Model:
    """
    Public SEN2SR model interface for HuggingFace Hub users.

    Parameters
    ----------
    local_dir : str
        Path to the directory returned by ``snapshot_download(repo_id)``.
    **overrides
        Optionally override any value from config.json, e.g.
        ``Model(local_dir=d, patch_size=256, overlap=64)``.
    """

    def __init__(self, local_dir: str, **overrides):
        config_path = os.path.join(local_dir, "config.json")
        with open(config_path) as f:
            config = json.load(f)

        config.update(overrides)

        if local_dir not in sys.path:
            sys.path.insert(0, local_dir)

        sen2sr_pt   = _load_module("sen2sr_pt", os.path.join(local_dir, "sen2sr_pt.py"))
        self._model = sen2sr_pt.SEN2SRPT(local_dir=local_dir, config=config)
        self.description = config.get("description", "")

    def predict(self, image: np.ndarray) -> np.ndarray:
        """
        Run 4x super-resolution on a single image.

        Parameters
        ----------
        image : (C, H, W) float32 numpy array, values in [0, 1]
                C must equal in_channels (4 for RGBN)

        Returns
        -------
        (C, H*4, W*4) float32 numpy array
        """
        return self._model.predict(image)

    def predict_tif(
        self,
        input_path:  str,
        output_path: str,
        bands:       Optional[List[int]] = None,
    ) -> None:
        """
        Full GeoTIFF super-resolution pipeline.

        Parameters
        ----------
        input_path  : path to input Sentinel-2 GeoTIFF
        output_path : output path for the 2.5 m SR GeoTIFF
        bands       : 0-based band indices to read (default: [0, 1, 2, 3])
        """
        self._model.predict_tif(input_path, output_path, bands)