import numba
import tempfile
numba.core.config.CACHE_DIR = tempfile.mkdtemp()
import json
import numpy as np
import pandas as pd
import torch
import seqpro
import streamlit as st
from model.module import CNN_Encoder_Model


def device_detect() -> tuple[torch.device, type]:
    """
    Detect whether the device is a CPU or GPU.
    Returns a tuple of (device, Tensor type)
    """
    if torch.cuda.is_available():
        device = torch.device("cuda")
        Tensor = torch.cuda.FloatTensor
    else:
        device = torch.device("cpu")
        Tensor = torch.FloatTensor
    return device, Tensor


def model_loading(config_local: str, epoch_local: str) -> CNN_Encoder_Model:
    """
    Load model parameters and return the model in evaluation mode.
    """
    device, _ = device_detect()
    with open(config_local, "r") as file:
        model_config = json.load(file)
        INPUTSHAPE = {
            "gene": 1,
            "reg": (4, 5000),
            "cds": (4, 4500),
            "trans": (1024,),
            "rna": (1,),
            "expression": [1],
        }
    model = CNN_Encoder_Model(model_config=model_config, input_shapes=INPUTSHAPE)
    model.to(device)
    checkpoint = torch.load(epoch_local, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return model


def format_sequence(utr_5, cds, utr_3) -> tuple[np.ndarray, np.ndarray]:
    utr_formatted = list()
    cds_formatted = list()
    for i in range(len(utr_5)):
        cds_formatted.append(cds[i][:4500].ljust(4500, "N"))
        utr_formatted.append(
            utr_5[i][-1000:].rjust(1000, "N") + utr_3[i][-4000:].ljust(4000, "N")
        )
    return seqpro.ohe(
        np.array(utr_formatted), alphabet=seqpro.alphabets.DNA
    ), seqpro.ohe(np.array(cds_formatted), alphabet=seqpro.alphabets.DNA)


def format_mrna_env(envs_local: str):
    mrna_env = pd.read_csv(envs_local, header=None, names=["value"])["value"].to_numpy()
    return mrna_env


def predict(
    model,
    utr: np.ndarray,
    cds: np.ndarray,
    env: np.ndarray,
    rna_count=6.3,
    batch_size=256,
) -> np.ndarray:
    """
    Predicate by model.
    """
    device, Tensor = device_detect()
    res = []
    with torch.no_grad():
        for line in range(0, len(utr), batch_size):
            count = utr[line : line + batch_size].shape[0]
            out = (
                model(
                    Tensor(utr[line : line + batch_size])
                    .to(device)
                    .transpose(2, 1)
                    .float(),
                    Tensor(cds[line : line + batch_size])
                    .to(device)
                    .transpose(2, 1)
                    .float(),
                    Tensor(env).to(device).repeat(count, 1),
                    Tensor([rna_count]).to(device).reshape(1, 1).repeat(count, 1),
                )[:, 0]
                .to("cpu")
                .detach()
                .numpy()
            )
            res.append(out)
    out = np.concatenate(res).reshape(-1)
   
    return out