# Custom Rust modules
import mol_fingerprint as mf
import bitexpand

import pandas as pd
import numpy as np
from pathlib import Path
from argparse import ArgumentParser

import zarr
from zarr.storage import LocalStore

from sklearn.preprocessing import normalize
from parallelbar import progress_map

import onnxruntime as ort
print('CUDA available:', 'CUDAExecutionProvider' in ort.get_available_providers())



def parse_args():
    parser = ArgumentParser()
    parser.add_argument("-i","--input-tsv", type=Path, required=True, help="Input TSV file with molecules to screen. It must not have any labels and has only two columns, 1st molecule_label, 2nd molecule_smiles")
    parser.add_argument("-z","--zarr-protein", type=Path, required=True, help="Input Zarr DB that holds the protein ids and their normalized projection vectors")
    return parser

parser = parse_args()

args = parser.parse_args()

# path_fp_basedir : Path = args.input_dir
# my_id :int = args.chunk_id

path_metabolites : Path = args.input_tsv
if not path_metabolites.is_file():
    raise ValueError(f"{path_metabolites} IS MISSING OR NOT A FILE")

df_metabolite = pd.read_csv(
    path_metabolites,
    sep="\t",
    header=None,
    names=["smiles" ,"compound_id"]
)

path_zarr = args.zarr_protein
store_zarr = LocalStore(path_zarr , read_only=True)
group_zarr = zarr.open(store_zarr, mode="r")
protein_vecs = (group_zarr["normalized"][:])


path_onnx = Path("/app/drug_projector.onnx")

def load_model(model_path):
    """Load ONNX model with GPU support"""
    model_path = Path(model_path)
    
    # Set up providers for GPU execution
    # providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    providers = ['CPUExecutionProvider']
    
    # Create inference session
    session = ort.InferenceSession(str(model_path), providers=providers)
    
    # Verify GPU is being used
    print(f"Available providers: {session.get_providers()}")
    
    return session

session = load_model(path_onnx)

def run_inference(input_vector):
    """Run inference on the model"""
    # Get input/output names
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    
    # Ensure input is the correct shape and type
    # if input_vector.shape != (2048,):
    #     raise ValueError(f"Input shape should be (2048,), got {input_vector.shape}")
    
    # Add batch dimension if needed
    input_data = input_vector
    
    # Run inference
    outputs = session.run([output_name], {input_name: input_data})
    
    # Return the output vector (remove batch dimension)
    return outputs[0]

def normalize_vectors_sklearn(vectors):
    """Normalize using sklearn - often fastest for large arrays"""
    return normalize(vectors, norm='l2', axis=1)

def project_drugs(smiles : list[str]) -> np.ndarray :
    fps = bitexpand.expand_bits(np.array([
        np.frombuffer(fp,dtype=np.uint8)
        for fp in mf.generate_fingerprints(smiles)
    ]))
    return normalize_vectors_sklearn(run_inference(fps))

pvecs = project_drugs(df_metabolite["smiles"].values.tolist())

my_results = np.dot(protein_vecs,pvecs.T)

pids = group_zarr["ids"][:]

path_outdir = Path("./")
def write_protein(i : int):
    df_out = df_metabolite[["compound_id"]].copy()
    df_out["protein_id"] = pids[i]
    df_out["score"] = my_results[i]
    # df_out["protein_sequence"] = group_zarr["sequences"][i]
    path_out = path_outdir / f"{pids[i]}_results.tsv"
    
    df_out.to_csv(
        path_out,
        sep="\t",
        index=False,
        header=False
    )

tasks = list(range(my_results.shape[0]))
progress_map(
    write_protein, tasks
)

print("PROCESS COMPLETED SUCCESSFULLY")