# 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 zarr.codecs import ZstdCodec

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_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
    
    # Run inference
    outputs = session.run([output_name], {input_name: input_vector})
    
    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"][:]

# Create output zarr store
path_outdir = Path("./")
zarr_name = path_zarr.stem + "_results.zarr"
path_output_zarr = path_outdir / zarr_name

# Create zarr store
store_output = LocalStore(path_output_zarr, read_only=False)
group_output = zarr.open(store_output, mode="w")

out_compress =[
    ZstdCodec( level=13 )
]

# Store the data
# Similarity matrix: (n_proteins, n_compounds)
z1 = group_output.create_array(
    "scores",
    dtype="float32",
    shape=my_results.shape,
    chunks=(min(1000, my_results.shape[0]), my_results.shape[1]),
    compressors=out_compress
)
z1[:] = my_results

# Compound IDs
z2 = group_output.create_array(
    "compound_ids",
    dtype="str",
    shape=(len(df_metabolite),),
    compressors=out_compress
)
z2[:] = df_metabolite["compound_id"].values

# Protein IDs
z3 = group_output.create_array(
    "protein_ids",
    dtype="str",
    shape=(len(pids),),
    compressors=out_compress
)
z3[:] = pids

# Store metadata
group_output.attrs['n_proteins'] = my_results.shape[0]
group_output.attrs['n_compounds'] = my_results.shape[1]
group_output.attrs['source_zarr'] = str(path_zarr)

print(f"Results saved to {path_output_zarr}")
print(f"Shape: {my_results.shape[0]} proteins x {my_results.shape[1]} compounds")
print("PROCESS COMPLETED SUCCESSFULLY")