"""
Needed objects for tests
"""

#---------------------------------------------------------------------------------------
# Dependencies
import pytest
import pandas as pd
import pickle
import numpy as np
import torch

from data_preprocessing.create_descriptors  import create_cleaned_mol_objects

#---------------------------------------------------------------------------------------
# Define fixtures

#---------------------------------------------------------------------------------------
# Data preprocessing
@pytest.fixture(scope="session")
def input_molecule_formats():
    class Formats:
        smiles = "CCO"
        smiles_coma = "CCO, CCN"
        smiles_list = ["CCO", "CCN"]
        smiles_df = pd.DataFrame({"smiLES": ["CCO", "CCN"]})
        smiles_df_wrong_key = pd.DataFrame({"notSMILES": ["CCO", "CCN"]})
    return Formats()

@pytest.fixture(scope="session")
def input_smiles():
    current_loc = __file__.rsplit("/",3)[0]
    with open(current_loc + "/assets/test_reference_data/smiles.pkl", "rb") as fl:
        input_smiles = pickle.load(fl)
    return input_smiles

@pytest.fixture(scope="session")
def input_mols_from_smiles():
    current_loc = __file__.rsplit("/",3)[0]
    with open(current_loc + "/assets/test_reference_data/smiles.pkl", "rb") as fl:
        input_smiles = pickle.load(fl)
    
    input_molecules = create_cleaned_mol_objects(input_smiles)
    return input_molecules

@pytest.fixture(scope="session")
def ecfps_from_smiles():
    current_loc = __file__.rsplit("/",3)[0]
    ecfps = np.load(current_loc + "/assets/test_reference_data/ecfps.npy")
    return ecfps

@pytest.fixture(scope="session")
def rdkit_descrs_from_smiles():
    current_loc = __file__.rsplit("/",3)[0]
    rdkit_descrs = np.load(current_loc + "/assets/test_reference_data/rdkit_descrs.npy")
    return rdkit_descrs

@pytest.fixture(scope="session")
def rdkit_descr_quantils():
    current_loc = __file__.rsplit("/",3)[0]
    rdkit_descr_quantils = np.load(
        current_loc + "/assets/test_reference_data/rdkit_descr_quantils.npy")
    return rdkit_descr_quantils

@pytest.fixture(scope="session")
def preprocessed_features():
    current_loc = __file__.rsplit("/",3)[0]
    preprocessed_features = np.load(
        current_loc + "/assets/test_reference_data/preprocessed_features.npy")
    return preprocessed_features

#---------------------------------------------------------------------------------------
# Model
@pytest.fixture(scope="session")
def model_input_query():
    current_loc = __file__.rsplit("/",3)[0]
    model_input_query = torch.load(
        current_loc + "/assets/test_reference_data/model_input_query.pt")
    return model_input_query

@pytest.fixture(scope="session")
def model_input_support_actives():
    current_loc = __file__.rsplit("/",3)[0]
    model_input_support_actives = torch.load(
        current_loc + "/assets/test_reference_data/model_input_support_actives.pt")
    return model_input_support_actives

@pytest.fixture(scope="session")
def model_input_support_inactives():
    current_loc = __file__.rsplit("/",3)[0]
    model_input_support_inactives = torch.load(
        current_loc + "/assets/test_reference_data/model_input_support_inactives.pt")
    return model_input_support_inactives

@pytest.fixture(scope="session")
def model_predictions():
    current_loc = __file__.rsplit("/",3)[0]
    model_predictions = torch.load(
        current_loc + "/assets/test_reference_data/model_predictions.pt")
    return model_predictions