# geney/tis/model.py — TITER model configuration, architecture, and loading.
"""
Translation Initiation Site (TIS) prediction using TITER.

TITER (Translation Initiation siTE detectoR) is a deep learning model that predicts
whether a codon is likely to be used as a translation initiation site.

Reference: https://github.com/zhangsaithu/titer
"""
from __future__ import annotations

import logging
import os
from typing import Dict

import numpy as np

logger = logging.getLogger(__name__)

import torch
import torch.nn as nn


# Lazy-loaded containers
_titer_model = None
_titer_model_dir = None
_codon_prior = None

# Number of ensemble models
N_ENSEMBLE_MODELS = 32

# Lazy-loaded PyTorch models
_titer_models_pt = None
_titer_device = None


# ============================================================================
# Configuration
# ============================================================================

def set_titer_model_dir(path: str) -> None:
    """
    Set the directory containing TITER model weights and codon prior.

    Expected files:
        - bestmodel_0.hdf5 ... bestmodel_31.hdf5 (32 ensemble models)
        - dict_piror_front_Gaotrain.npy (codon prior, optional but recommended)
    """
    global _titer_model_dir, _titer_model, _codon_prior
    if not os.path.isdir(path):
        raise ValueError(f"TITER model directory does not exist: {path}")
    _titer_model_dir = path
    _titer_model = None  # Reset cached model
    _codon_prior = None  # Reset cached prior


def get_titer_model_dir() -> str:
    """Get the TITER model directory from environment or raise an error."""
    global _titer_model_dir
    if _titer_model_dir is not None:
        return _titer_model_dir

    env_path = os.environ.get('TITER_MODEL_DIR')
    if env_path and os.path.isdir(env_path):
        _titer_model_dir = env_path
        return _titer_model_dir

    # Auto-detect bundled models at geney/models/titer/
    pkg_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    default = os.path.join(pkg_dir, 'models', 'titer')
    if os.path.isdir(default):
        _titer_model_dir = default
        return _titer_model_dir

    raise ValueError(
        "TITER model directory not set. Either:\n"
        "  1. Call set_titer_model_dir('/path/to/titer/model/')\n"
        "  2. Set TITER_MODEL_DIR environment variable\n"
        "Models can be downloaded from: https://github.com/zhangsaithu/titer\n"
        "Required: bestmodel_0.hdf5 ... bestmodel_31.hdf5\n"
        "Optional: dict_piror_front_Gaotrain.npy (codon prior)"
    )


# ============================================================================
# Codon Prior (CodonScore)
# ============================================================================

def _load_codon_prior() -> Dict[str, float]:
    """
    Load the codon prior dictionary.

    The prior represents the preference of codon composition at TIS sites,
    calculated as: log+(#codon at TIS / #TIS) / (#codon in background / #background)

    Returns dict mapping codon -> prior score.
    """
    global _codon_prior

    if _codon_prior is not None:
        return _codon_prior

    model_dir = get_titer_model_dir()
    prior_path = os.path.join(model_dir, "dict_piror_front_Gaotrain.npy")

    if os.path.exists(prior_path):
        try:
            loaded = np.load(prior_path, allow_pickle=True)
            # Handle numpy 0-d array containing dict
            if loaded.ndim == 0:
                _codon_prior = loaded.item()
            else:
                _codon_prior = dict(loaded)
            logger.debug("Loaded codon prior with %d codons", len(_codon_prior))
        except Exception as e:
            logger.warning("Could not load codon prior: %s", e)
            _codon_prior = _get_default_codon_prior()
    else:
        logger.warning("Codon prior not found at %s, using defaults", prior_path)
        _codon_prior = _get_default_codon_prior()

    return _codon_prior


def _get_default_codon_prior() -> Dict[str, float]:
    """
    Default codon prior based on paper statistics.

    Top codons: AUG > CUG > GUG > UUG > ACG
    (Using DNA notation: ATG, CTG, GTG, TTG, ACG)
    """
    # Based on Figure 2a and Supplementary Fig S1 from the paper
    # These are approximate values - the real prior should be loaded from file
    return {
        'ATG': 5.0,   # AUG - dominant (~50%)
        'CTG': 2.5,   # CUG - second most common
        'GTG': 2.0,   # GUG
        'TTG': 1.8,   # UUG (TTG in DNA)
        'ACG': 1.5,   # ACG
        'AAG': 1.0,
        'AGG': 1.0,
        'ATA': 0.8,
        'ATC': 0.8,
        'ATT': 0.8,
        # Other codons get low prior
    }


def get_codon_score(codon: str) -> float:
    """
    Get the CodonScore (prior) for a given codon.

    From the paper: CodonScore represents the preference of codon composition
    at translation initiation sites.

    Args:
        codon: 3-letter codon (DNA notation, e.g., 'ATG')

    Returns:
        CodonScore value (higher = more likely to be TIS)
    """
    prior = _load_codon_prior()
    codon = codon.upper().replace('U', 'T')  # Convert RNA to DNA

    score = prior.get(codon, 0.0)

    # Handle 'never' entries from the original prior
    if score == 'never' or score < 0:
        return 0.0

    return float(score)


# ============================================================================
# Model Architecture and Loading (PyTorch)
# ============================================================================

def set_titer_device(device: str) -> None:
    """Override the device used by TITER models (e.g. 'cpu', 'cuda', 'mps')."""
    global _titer_device, _titer_models_pt
    _titer_device = torch.device(device)
    _titer_models_pt = None  # Force reload on next use


def get_titer_device():
    """Return the current TITER device (auto-detected if not overridden)."""
    return _get_torch_device()


def _get_torch_device():
    """Get the best available device for PyTorch."""
    global _titer_device
    if _titer_device is not None:
        return _titer_device

    import sys
    if sys.platform == 'darwin' and torch.backends.mps.is_available():
        try:
            torch.tensor([1.0], device="mps")
            _titer_device = torch.device("mps")
        except RuntimeError:
            _titer_device = torch.device("cpu")
    elif torch.cuda.is_available():
        _titer_device = torch.device("cuda")
    else:
        _titer_device = torch.device("cpu")
    return _titer_device


class TITERModel(nn.Module):
    """
    TITER model architecture in PyTorch.

    Architecture:
        - Conv1D: 128 filters, kernel_size=3, ReLU
        - MaxPool1D: pool_size=3
        - Dropout: 0.2137
        - LSTM: 256 units, return_sequences=True
        - Dropout: 0.7238
        - Flatten
        - Linear: 1 output
        - Sigmoid
    """

    def __init__(self):
        super(TITERModel, self).__init__()

        # Conv1D: input (batch, seq_len=203, channels=8) -> (batch, seq_len-2, 128)
        # Keras Conv1D: (batch, steps, features) with kernel on steps axis
        # PyTorch Conv1d: (batch, channels, seq_len) - need to transpose
        self.conv1d = nn.Conv1d(
            in_channels=8,
            out_channels=128,
            kernel_size=3,
            padding=0  # 'valid' padding
        )
        self.relu = nn.ReLU()

        # MaxPool1D: pool_size=3
        self.maxpool = nn.MaxPool1d(kernel_size=3)

        self.dropout1 = nn.Dropout(p=0.21370950078747658)

        # LSTM: input_size=128, hidden_size=256, batch_first=True
        # After conv+pool: seq_len = (203-2) // 3 = 67
        self.lstm = nn.LSTM(
            input_size=128,
            hidden_size=256,
            num_layers=1,
            batch_first=True,
            bidirectional=False
        )

        self.dropout2 = nn.Dropout(p=0.7238091317104384)

        # Flatten: 67 * 256 = 17152
        # But actually: (203-2)/3 = 67 (floor), so 67*256 = 17152
        self.fc = nn.Linear(67 * 256, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x shape: (batch, 203, 8)

        # Transpose for Conv1d: (batch, 8, 203)
        x = x.transpose(1, 2)

        # Conv1d + ReLU: (batch, 128, 201)
        x = self.relu(self.conv1d(x))

        # MaxPool: (batch, 128, 67)
        x = self.maxpool(x)

        x = self.dropout1(x)

        # Transpose back for LSTM: (batch, 67, 128)
        x = x.transpose(1, 2)

        # LSTM: (batch, 67, 256)
        x, _ = self.lstm(x)

        x = self.dropout2(x)

        # Flatten: (batch, 67*256)
        x = x.reshape(x.size(0), -1)

        # FC + Sigmoid: (batch, 1)
        x = self.sigmoid(self.fc(x))

        return x


def convert_keras_to_pytorch(keras_weights_path: str, pytorch_model: TITERModel) -> TITERModel:
    """
    Convert Keras HDF5 weights to PyTorch model.

    Args:
        keras_weights_path: Path to .hdf5 file
        pytorch_model: PyTorch TITERModel instance

    Returns:
        PyTorch model with loaded weights
    """
    import h5py

    def find_layer_group(f, layer_prefix):
        """Find the group for a layer (handles varying suffixes like conv1d_1, conv1d_3, etc.)"""
        model_weights = f['model_weights']
        for key in model_weights.keys():
            if key.startswith(layer_prefix):
                # Found it - return the innermost group (layer_name/layer_name/)
                return model_weights[key][key]
        raise KeyError(f"Could not find layer starting with '{layer_prefix}' in {list(model_weights.keys())}")

    with h5py.File(keras_weights_path, 'r') as f:
        # Load Conv1D weights
        # Keras Conv1D kernel shape: (kernel_size, in_channels, out_channels) = (3, 8, 128)
        # PyTorch Conv1d weight shape: (out_channels, in_channels, kernel_size) = (128, 8, 3)
        conv_group = find_layer_group(f, 'conv1d')
        conv_kernel = np.array(conv_group['kernel:0'])
        conv_bias = np.array(conv_group['bias:0'])

        # Transpose: (3, 8, 128) -> (128, 8, 3)
        conv_kernel = np.transpose(conv_kernel, (2, 1, 0))
        pytorch_model.conv1d.weight.data = torch.from_numpy(conv_kernel.copy()).float()
        pytorch_model.conv1d.bias.data = torch.from_numpy(conv_bias.copy()).float()

        # Load LSTM weights
        # Keras LSTM:
        #   kernel: (input_size=128, 4*hidden_size=1024) - input weights
        #   recurrent_kernel: (hidden_size=256, 4*hidden_size=1024) - hidden weights
        #   bias: (4*hidden_size=1024,) - single bias
        # PyTorch LSTM:
        #   weight_ih_l0: (4*hidden_size, input_size)
        #   weight_hh_l0: (4*hidden_size, hidden_size)
        #   bias_ih_l0: (4*hidden_size,)
        #   bias_hh_l0: (4*hidden_size,)
        lstm_group = find_layer_group(f, 'lstm')
        lstm_kernel = np.array(lstm_group['kernel:0'])
        lstm_recurrent = np.array(lstm_group['recurrent_kernel:0'])
        lstm_bias = np.array(lstm_group['bias:0'])

        # Keras LSTM gate order: i, f, c, o (same as PyTorch)
        # Transpose kernel: (128, 1024) -> (1024, 128)
        pytorch_model.lstm.weight_ih_l0.data = torch.from_numpy(lstm_kernel.T.copy()).float()

        # Transpose recurrent: (256, 1024) -> (1024, 256)
        pytorch_model.lstm.weight_hh_l0.data = torch.from_numpy(lstm_recurrent.T.copy()).float()

        # Keras uses single bias, PyTorch has two biases
        # Put all bias in bias_ih, set bias_hh to zeros
        pytorch_model.lstm.bias_ih_l0.data = torch.from_numpy(lstm_bias.copy()).float()
        pytorch_model.lstm.bias_hh_l0.data = torch.zeros(1024).float()

        # Load Dense weights
        # Keras Dense kernel shape: (in_features=17152, out_features=1)
        # PyTorch Linear weight shape: (out_features=1, in_features=17152)
        dense_group = find_layer_group(f, 'dense')
        dense_kernel = np.array(dense_group['kernel:0'])
        dense_bias = np.array(dense_group['bias:0'])

        # Transpose: (17152, 1) -> (1, 17152)
        pytorch_model.fc.weight.data = torch.from_numpy(dense_kernel.T.copy()).float()
        pytorch_model.fc.bias.data = torch.from_numpy(dense_bias.copy()).float()

    return pytorch_model


def _load_titer_models_pytorch():
    """Load all TITER ensemble models in PyTorch."""
    global _titer_models_pt

    if _titer_models_pt is not None:
        return _titer_models_pt

    model_dir = get_titer_model_dir()
    device = _get_torch_device()

    logger.info("Loading TITER models to %s...", device)

    _titer_models_pt = []

    for i in range(N_ENSEMBLE_MODELS):
        # Check for PyTorch format first
        pt_path = os.path.join(model_dir, f"bestmodel_{i}.pt")
        hdf5_path = os.path.join(model_dir, f"bestmodel_{i}.hdf5")

        model = TITERModel()

        if os.path.exists(pt_path):
            # Load PyTorch weights directly
            state_dict = torch.load(pt_path, map_location=device, weights_only=True)
            model.load_state_dict(state_dict)
        elif os.path.exists(hdf5_path):
            # Convert from Keras
            model = convert_keras_to_pytorch(hdf5_path, model)
        else:
            raise FileNotFoundError(
                f"TITER model file not found: {pt_path} or {hdf5_path}\n"
                f"Expected {N_ENSEMBLE_MODELS} model files in {model_dir}"
            )

        model = model.to(device)
        model.eval()
        _titer_models_pt.append(model)

    logger.info("TITER loaded (%d models)", N_ENSEMBLE_MODELS)
    return _titer_models_pt


def convert_all_models_to_pytorch(model_dir: str = None, save: bool = True):
    """
    Convert all Keras HDF5 models to PyTorch format.

    Args:
        model_dir: Directory containing .hdf5 files (default: TITER model dir)
        save: If True, save .pt files alongside .hdf5 files

    Returns:
        List of converted PyTorch models
    """
    if model_dir is None:
        model_dir = get_titer_model_dir()

    models = []
    for i in range(N_ENSEMBLE_MODELS):
        hdf5_path = os.path.join(model_dir, f"bestmodel_{i}.hdf5")
        pt_path = os.path.join(model_dir, f"bestmodel_{i}.pt")

        if not os.path.exists(hdf5_path):
            raise FileNotFoundError(f"Keras model not found: {hdf5_path}")

        logger.info("Converting model %d...", i)
        model = TITERModel()
        model = convert_keras_to_pytorch(hdf5_path, model)
        models.append(model)

        if save:
            torch.save(model.state_dict(), pt_path)
            logger.info("  Saved: %s", pt_path)

    logger.info("Conversion complete. %d models converted.", N_ENSEMBLE_MODELS)
    return models


# Keep old function name for compatibility but use PyTorch
def _load_titer_model():
    """Lazy load the TITER models (PyTorch version)."""
    return _load_titer_models_pytorch()
