# (c) 2023 Sormannilab and Aubin Ramon
#
# Lightning testing of the AbNatiV model.
#
# ============================================================================

from .model.scoring_functions import abnativ_scoring, abnativ_scoring_csv
from .model.onehotencoder import alphabet
import argparse
import os

from .model.utils import is_protein
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq

from .init import PRETRAINED_MODELS_DIR

def run(args: argparse.Namespace):

    # Check that models are downloaded
    if not os.path.exists(PRETRAINED_MODELS_DIR):
        raise Exception("Models not found. Please run 'abnativ update' to download the models.")

    if not os.path.exists(args.output_directory):
        os.makedirs(args.output_directory)

    ## DATA SCORING ##
    batch_size = 128

    # Check if input is a CSV file
    is_csv_input = (
        os.path.isfile(args.input_filepath_or_seq) and
        args.input_filepath_or_seq.lower().endswith('.csv')
    )

    if is_csv_input:
        # Use CSV scoring with column preservation
        abnativ_df_mean, abnativ_df_profile = abnativ_scoring_csv(
            model_type=args.nativeness_type,
            input_csv=args.input_filepath_or_seq,
            col_id=args.col_id,
            col_seq=args.col_seq,
            batch_size=batch_size,
            chunk_size=args.chunk_size,
            mean_score_only=args.mean_score_only,
            do_align=args.do_align,
            is_VHH=args.is_VHH,
            is_plotting_profiles=args.is_plotting_profiles,
            output_dir=args.output_directory,
            output_id=args.output_id,
            verbose=True,
            run_parall_al=args.ncpu,
            preserve_columns=True,
        )

    else:
        # Original FASTA/sequence handling
        # Load the fasta file
        if os.path.isfile(args.input_filepath_or_seq):
            seq_records = list(SeqIO.parse(args.input_filepath_or_seq, 'fasta'))

        # If just a single sequence string
        elif is_protein(args.input_filepath_or_seq, alphabet + ['X']):
            seq_records = [SeqRecord(Seq(args.input_filepath_or_seq), id='single_seq')]

        else:
            raise Exception(f'Can not find the input file {args.input_filepath_or_seq} or if you gave a single protein sequence make sure it is composed only of {alphabet}')

        abnativ_df_mean, abnativ_df_profile = abnativ_scoring(
            args.nativeness_type, seq_records, batch_size, args.mean_score_only,
            args.do_align, args.is_VHH, args.is_plotting_profiles, args.output_directory, args.output_id,
            run_parall_al=args.ncpu
        )

    ## DATA SAVING ##
    print(f'\n-> Scores being saved in {args.output_directory}\n')
    abnativ_df_mean.to_csv(os.path.join(args.output_directory, f'{args.output_id}_abnativ_seq_scores.csv'), index=False)
    if not args.mean_score_only:
        abnativ_df_profile.to_csv(os.path.join(args.output_directory, f'{args.output_id}_abnativ_res_scores.csv'), index=False)
    if args.is_plotting_profiles:
        save_profile_fp = os.path.join(args.output_directory, f'{args.output_id}_profiles')
        print(f'\n-> Profile plots saved in {save_profile_fp}\n')
