"""Process mutation files"""
import logging
import os
import shutil
import subprocess
import tempfile

import pandas as pd
import synapseclient
from synapseclient import Synapse
from synapseclient.core.exceptions import SynapseTimeoutError

from . import process_functions

logger = logging.getLogger(__name__)

# Some columns are already capitalized, so they aren't included here
MAF_COL_MAPPING = {
    'HUGO_SYMBOL': 'Hugo_Symbol',
    'ENTREZ_GENE_ID': 'Entrez_Gene_Id',
    'CENTER': 'Center',
    'NCBI_BUILD': 'NCBI_Build',
    'CHROMOSOME': 'Chromosome',
    'START_POSITION': 'Start_Position',
    'END_POSITION': 'End_Position',
    'STRAND': 'Strand',
    'VARIANT_CLASSIFICATION': 'Variant_Classification',
    'VARIANT_TYPE': 'Variant_Type',
    'REFERENCE_ALLELE': 'Reference_Allele',
    'TUMOR_SEQ_ALLELE1': 'Tumor_Seq_Allele1',
    'TUMOR_SEQ_ALLELE2': 'Tumor_Seq_Allele2',
    'DBSNP_RS': 'dbSNP_RS',
    'DBSNP_VAL_STATUS': 'dbSNP_Val_Status',
    'TUMOR_SAMPLE_BARCODE': 'Tumor_Sample_Barcode',
    'MATCHED_NORM_SAMPLE_BARCODE': 'Matched_Norm_Sample_Barcode',
    'MATCH_NORM_SEQ_ALLELE1': 'Match_Norm_Seq_Allele1',
    'MATCH_NORM_SEQ_ALLELE2': 'Match_Norm_Seq_Allele2',
    'TUMOR_VALIDATION_ALLELE1': 'Tumor_Validation_Allele1',
    'TUMOR_VALIDATION_ALLELE2': 'Tumor_Validation_Allele2',
    'MATCH_NORM_VALIDATION_ALLELE1': 'Match_Norm_Validation_Allele1',
    'MATCH_NORM_VALIDATION_ALLELE2': 'Match_Norm_Validation_Allele2',
    'VERIFICATION_STATUS': 'Verification_Status',
    'VALIDATION_STATUS': 'Validation_Status',
    'MUTATION_STATUS': 'Mutation_Status',
    'SEQUENCING_PHASE': 'Sequencing_Phase',
    'SEQUENCE_SOURCE': 'Sequence_Source',
    'VALIDATION_METHOD': 'Validation_Method',
    'SCORE': 'Score',
    'BAM_FILE': 'BAM_File',
    'SEQUENCER': 'Sequencer',
    'T_REF_COUNT': 't_ref_count',
    'T_ALT_COUNT': 't_alt_count',
    'N_REF_COUNT': 'n_ref_count',
    'N_ALT_COUNT': 'n_alt_count',
    'ALLELE': 'Allele',
    'AMINO_ACID_CHANGE': 'amino_acid_change',
    'AMINO_ACIDS': 'Amino_acids',
    'CDS_POSITION': 'CDS_position',
    'CODONS': 'Codons',
    'CONSEQUENCE': 'Consequence',
    'EXISTING_VARIATION': 'Existing_variation',
    'EXON_NUMBER': 'Exon_Number',
    'FEATURE': 'Feature',
    'FEATURE_TYPE': 'Feature_type',
    'GENE': 'Gene',
    'HGVSC': 'HGVSc',
    'HGVSP': 'HGVSp',
    'HGVSP_SHORT': 'HGVSp_Short',
    'HOTSPOT': 'Hotspot',
    'MA:FIMPACT': 'MA:FImpact',
    'MA:LINK.MSA': 'MA:link.MSA',
    'MA:LINK.PDB': 'MA:link.PDB',
    'MA:LINK.VAR': 'MA:link.var',
    'MA:PROTEIN.CHANGE': 'MA:protein.change',
    'POLYPHEN': 'PolyPhen',
    'PROTEIN_POSITION': 'Protein_position',
    'REFSEQ': 'RefSeq',
    'TRANSCRIPT': 'transcript',
    'TRANSCRIPT_ID': 'Transcript_ID',
    'ALL_EFFECTS': 'all_effects',
    'CDNA_CHANGE': 'cdna_change',
    'CDNA_POSITION': 'cDNA_position',
    'N_DEPTH': 'n_depth',
    'T_DEPTH': 't_depth'
}

KNOWN_STRING_COLS = ['IS_NEW', 'ALLELE_NUM', 'Chromosome', 'CLIN_SIG',
                     'MOTIF_NAME', 'HIGH_INF_POS', 'MINIMISED',
                     'CHROMOSOME', 'VERIFICATION_STATUS',
                     'VALIDATION_STATUS', 'MUTATION_STATUS',
                     'SEQUENCE_SOURCE', 'SEQUENCER', 'REPORT_AF',
                     'CDNA_CHANGE', 'AMINO_ACID_CHANGE',
                     'TRANSCRIPT', 'STRAND_VEP',
                     'HGNC_ID', 'PUBMED', 'PICK']


def _rename_column_headers(dataframe: pd.DataFrame,
                           col_map: dict) -> pd.DataFrame:
    """Rename dataframe column headers

    Args:
        dataframe: pandas dataframe
        col_map: Column Mapping {column_name: new_column_name,}

    Returns:
        Dataframe with new columns

    """
    dataframe = dataframe.rename(columns=col_map)
    return dataframe


def _convert_to_str_dtype(column_types, known_string_cols):
    """Sometimes the deteremined dtype is incorrect based off the first
    100 rows, update the incorrect dtypes.
    """
    for str_col in known_string_cols:
        if column_types.get(str_col):
            column_types[str_col] = 'object'
    return column_types


def determine_dtype(path: str):
    """Reads in a dataframe partially and determines the dtype of columns"""
    subset_df = pd.read_csv(path, nrows=100, sep="\t")
    dtypes = subset_df.dtypes
    colnames = dtypes.index
    types = [i.name for i in dtypes.values]
    column_types = dict(zip(colnames, types))
    return column_types


def move_maf(mutation_path, input_files_dir):
    """Moves maf files into processing directory. Maf file's column headers
    are renamed if necessary"""
    header_df = pd.read_csv(mutation_path, sep="\t",
                            index_col=0, nrows=1, comment="#")
    # If any column headers need to be remapped, remap
    if sum(header_df.columns.isin(MAF_COL_MAPPING.keys())) > 0:
        filename = os.path.basename(mutation_path)
        column_types = determine_dtype(mutation_path)
        new_column_types = _convert_to_str_dtype(column_types,
                                                 KNOWN_STRING_COLS)
        mafdf = pd.read_csv(mutation_path, sep="\t", dtype=new_column_types)
        mafdf = _rename_column_headers(mafdf, col_map=MAF_COL_MAPPING)
        mafdf.to_csv(os.path.join(input_files_dir, filename),
                     sep="\t", index=False)
    else:
        shutil.copy(mutation_path, input_files_dir)


def move_mutation(mutation_path, input_files_dir):
    """Move mutation file into processing directory"""
    # If mutation file is vcf, just copy
    if mutation_path.endswith(".vcf"):
        shutil.copy(mutation_path, input_files_dir)
    else:
        move_maf(mutation_path, input_files_dir)


def process_mutation_workflow(syn: Synapse, center: str,
                              validfiles: pd.DataFrame,
                              genie_annotation_pkg: str,
                              database_mappingdf: str,
                              workdir: str) -> str:
    """Process vcf/maf workflow

    Args:
        syn: Synapse connection
        center: Center name
        validfiles: Center validated files
        genie_annotation_pkg: Genome Nexus annotation tools
        database_mappingdf: Database to synapse id mapping dataframe
        workdir: Working directory

    Returns:
        Annotated Maf Path

    """
    # Get valid files
    mutation_files = validfiles['fileType'].isin(["maf", "vcf"])
    valid_mutation_files = validfiles['path'][mutation_files].tolist()
    # If there are no valid mutation files, return
    if not valid_mutation_files:
        logger.info("No mutation data")
        return
    # Certificate to use GENIE Genome Nexus
    syn.get("syn22053204", ifcollision="overwrite.local",
            downloadLocation=genie_annotation_pkg)
    # Genome Nexus Jar file
    syn.get("syn22084320", ifcollision="overwrite.local",
            downloadLocation=genie_annotation_pkg)

    annotated_maf_path = annotate_mutation(
        center=center,
        mutation_files=valid_mutation_files,
        genie_annotation_pkg=genie_annotation_pkg,
        workdir=workdir
    )

    maf_tableid = database_mappingdf.Id[
        database_mappingdf['Database'] == 'vcf2maf'
    ].iloc[0]
    flatfiles_synid = database_mappingdf.Id[
        database_mappingdf['Database'] == "centerMaf"
    ].iloc[0]
    # Split into narrow maf and store into db / flat file
    split_and_store_maf(syn=syn,
                        center=center,
                        maf_tableid=maf_tableid,
                        annotated_maf_path=annotated_maf_path,
                        flatfiles_synid=flatfiles_synid,
                        workdir=workdir)

    return annotated_maf_path


def annotate_mutation(center: str, mutation_files: list,
                      genie_annotation_pkg: str,
                      workdir: str) -> str:
    """Process vcf/maf files

    Args:
        center: Center name
        mutation_files: list of mutation files
        genie_annotation_pkg: Path to GENIE annotation package

    Returns:
        Path to final maf
    """
    input_files_dir = tempfile.mkdtemp(dir=workdir)
    output_files_dir = tempfile.mkdtemp(dir=workdir)

    for mutation_file in mutation_files:
        move_mutation(mutation_file, input_files_dir)

    merged_maf_path = os.path.join(
        output_files_dir, f"data_mutations_extended_{center}.txt"
    )
    annotater_cmd = ['bash', os.path.join(genie_annotation_pkg,
                                          'annotation_suite_wrapper.sh'),
                     f'-i={input_files_dir}',
                     f'-o={output_files_dir}',
                     f'-m={merged_maf_path}',
                     f'-c={center}',
                     '-s=WXS',
                     f'-p={genie_annotation_pkg}']

    subprocess.check_call(annotater_cmd)

    return merged_maf_path


def append_or_createdf(dataframe: pd.DataFrame, filepath: str):
    """Creates a file with the dataframe or appends to a existing file.

    Args:
        df: pandas.dataframe to write out
        filepath: Filepath to append or create

    """
    if not os.path.exists(filepath) or os.stat(filepath).st_size == 0:
        dataframe.to_csv(filepath, sep="\t", index=False)
    else:
        dataframe.to_csv(filepath, sep="\t", mode='a', index=False,
                         header=None)


def store_full_maf(syn: Synapse, filepath: str, parentid: str):
    """Stores full maf file

    Args:
        syn: Synapse connection
        filepath: Path to file
        parentid: Synapse container id

    """
    syn.store(synapseclient.File(filepath, parentId=parentid))


def store_narrow_maf(syn: Synapse, filepath: str, maf_tableid: str):
    """
    Stores the narrow maf in Synapse Table

    Args:
        syn: Synapse connection
        filepath: Path to maf file
        maf_tableid: database synid

    """
    logger.info(f'STORING {filepath}')
    # database = syn.get(maf_tableid)
    try:
        update_table = synapseclient.Table(maf_tableid, filepath,
                                           separator="\t")
        syn.store(update_table)
    except SynapseTimeoutError:
        # This error occurs because of waiting for table to index.
        # Don't worry about this.
        pass


def format_maf(mafdf: pd.DataFrame, center: str) -> pd.DataFrame:
    """Format maf file, shortens the maf file length

    Args:
        mafdf: mutation dataframe
        center: Center name

    Returns:
        Formatted mutation dataframe"""
    mafdf['Center'] = center
    mafdf['Tumor_Sample_Barcode'] = [
        process_functions.checkGenieId(i, center)
        for i in mafdf['Tumor_Sample_Barcode']
    ]

    mafdf['Sequence_Source'] = float('nan')
    mafdf['Sequencer'] = float('nan')
    mafdf['Validation_Status'][
        mafdf['Validation_Status'].isin(["Unknown", "unknown"])
    ] = ''

    return mafdf


def split_and_store_maf(syn: 'Synapse', center: str, maf_tableid: str,
                        annotated_maf_path: str, flatfiles_synid: str,
                        workdir: str):
    """Separates annotated maf file into narrow and full maf and stores them

    Args:
        syn: Synapse connection
        center: Center
        maf_tableid: Mutation table synapse id
        annotated_maf_path: Annotated maf
        flatfiles_synid: GENIE flat files folder

    """
    narrow_maf_cols = [col['name']
                       for col in syn.getTableColumns(maf_tableid)
                       if col['name'] != 'inBED']
    full_maf_path = os.path.join(
        workdir, center, "staging",
        f"data_mutations_extended_{center}.txt"
    )
    narrow_maf_path = os.path.join(
        workdir, center, "staging",
        f"data_mutations_extended_{center}_MAF_narrow.txt"
    )
    maf_chunks = pd.read_csv(annotated_maf_path, sep="\t", chunksize=100000)

    for maf_chunk in maf_chunks:
        maf_chunk = format_maf(maf_chunk, center)
        append_or_createdf(maf_chunk, full_maf_path)
        narrow_maf_chunk = maf_chunk[narrow_maf_cols]
        append_or_createdf(narrow_maf_chunk, narrow_maf_path)

    store_narrow_maf(syn, narrow_maf_path, maf_tableid)
    # Store MAF flat file into synapse
    store_full_maf(syn, full_maf_path, flatfiles_synid)
