'''
Script to perform protein identifier remapping and various filtering steps
(e.g. intra-species and intra-pathogen ppi removal, redundancy,
onlt reviewed identifiers, etc.).

Input:

     - expects a merged PPI dataframe generated by the pathogen_selection.py
       script. In other words: a pandas generated .tsf file that conforms to
       the PSI-MITAB format and uses the following column headers:
       col_names = [
        'xref_A', 'xref_B', 'alt_identifiers_A', 'alt_identifiers_B',
        'aliases_A', 'aliases_B', 'detection_method', 'author', 'publication',
        'taxid_A', 'taxid_B', 'interaction_type', 'source_database_ids',
        'interaction_identifiers', 'confidence_score'
       ]

     - Also expects taxonid mapping files provided as a directory containing
       a mapping file for botht the pathogen and the associated/host taxonids.

The output is saved to a directory of choice (preferably the one mentioned
for the input files, which is named after the pathogen taxonid). The file names
will not overlap with the default ones used by the previous script.

NOTE: the local remapping option cannot discern between reviewed and unreviewed entries.
'''

import argparse
import numpy as np
import pandas as pd
import sys

from pathlib import Path

from phppipy.dataprep import taxonid
from phppipy.ppi_tools import id_mapper
from phppipy.ppi_tools import ppi_filter

parser = argparse.ArgumentParser(
    description='Script to filter and remap PPI datasets.',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
    '-i',
    '--input',
    dest='input',
    type=str,
    required=True,
    help='Path to merged and pathogen-specific PPI dataset.')
parser.add_argument(
    '-t',
    '--taxonid',
    dest='taxonid',
    type=str,
    required=True,
    help='Path to directory with taxonID lists.')
parser.add_argument(
    '-l',
    '--local',
    dest='local',
    type=str,
    required=False,
    help=
    'Path to file for local remapping. Leave this out to use online service.')
parser.add_argument(
    '-s',
    '--skip',
    action='store_true',
    help=
    'Allows user to skip the creation of mapping files and use existing ones.'
)
parser.add_argument(
    '-o',
    '--output',
    dest='output',
    type=str,
    required=True,
    help='Path to output directory.')
args = parser.parse_args()

# create output folder
out_path = Path(args.output)
# create new directory inside pathogen taxonid directory automatically?
# pathogen_taxonid = str(input_dir.parent)

# import ppi file
ppi_file = Path(args.input)
ppi_df = pd.read_csv(ppi_file, sep='\t', header=0)
print('PPIs were imported from {}\n'.format(ppi_file))

# import taxonID lists
input_dir = Path(args.taxonid)
pathogen_taxonid_path = [i for i in input_dir.glob('*child*')]
associated_taxonid_path = [i for i in input_dir.glob('*interacting*')]
# check if there is only 1 possible file match for the pathogen and host taxid files
try:
    assert len(pathogen_taxonid_path) == 1
    assert len(associated_taxonid_path) == 1
except AssertionError:
    print(
        'WARNING: multiple possible taxonID files were found in the directory {}'.
        format(input_dir))
    print(
        'Please run the pathogen_selection.py script again and verify that only two files are being created in a pathogen-specific directory, e.g. 10292/10292-child-taxids.txt and 10292/associated-taxids.txt'
    )
pathogen_taxonids = taxonid.read_taxids(pathogen_taxonid_path[0])
host_taxonids = taxonid.read_taxids(associated_taxonid_path[0])
print('TaxonIDs were read from {} and {}\n'.format(pathogen_taxonid_path[0],
                                                   associated_taxonid_path[0]))

# remove intra-species and intra-pathogen interactions
# temp_intra_species = ppi_df['inter-intra-species'].copy()
# temp_intra_pathogen = ppi_df['inter-intra-pathogen'].copy()
ppi_filter.annotate_inter_intra_species(ppi_df)
ppi_filter.annotate_inter_intra_pathogen(ppi_df, pathogen_taxonids)
# assert all(temp_intra_species == ppi_df['inter-intra-species'])
# assert all(temp_intra_pathogen == ppi_df['inter-intra-pathogen'])
ppi_df = ppi_df.loc[(ppi_df['inter-intra-species'] != 'intra')
                    & (ppi_df['inter-intra-pathogen'] != 'intra')]
ppi_df = ppi_df.reset_index(drop=True)
ppi_df.drop(
    ['inter-intra-species', 'inter-intra-pathogen'], axis=1, inplace=True)
print('All intra-species and intra-pathogen interactions were omitted.')
# check for intra-hosts cross species
assert any(
    ppi_df.taxid_A.isin(host_taxonids) & ppi_df.taxid_B.isin(host_taxonids)
) is False

# create unique identifier by combining xrefs
ppi_filter.unique_identifier(ppi_df)
ppi_df['original_unique'] = ~ppi_df.duplicated(subset=['xref_partners_sorted'])

# fix bad xref identifiers (i.e. containing more than a single entry)
id_mapper.check_unique_identifier(ppi_df)

# re-create unique identifier by combining xrefs
ppi_filter.unique_identifier(ppi_df)
ppi_df['checker_unique'] = ~ppi_df.duplicated(subset=['xref_partners_sorted'])
print(
    'By removing/fixing protein identifiers that consisted of multiple entries, the following duplicates were introduced in the dataset:\n'
)
print(ppi_df.loc[(ppi_df.original_unique)
                 & ~(ppi_df.checker_unique)].groupby('origin').size(), '\n')

# create UniProt mapping files and remap identifiers
try:
    out_mappings = out_path / 'mapping'
    out_mappings.mkdir(parents=True, exist_ok=False)
except FileExistsError:
    print(
        'Warning: supplied output directory already contains a "mapping" directory, aborting operation.'
    )
    sys.exit(1)
if args.local:
    full_mapping_file = Path(args.local)
    id_mapper.map2uniprot(
        ppi_df,
        out_mappings,
        reviewed_only=True,
        full_mapping_file=full_mapping_file,
        skip_creation=args.skip)
else:
    id_mapper.map2uniprot(ppi_df, out_mappings, reviewed_only=True, skip_creation=args.skip)

# remove multiple mappings
ppi_df = id_mapper.remove_mult(ppi_df)

# re-create unique identifier by combining xrefs
ppi_filter.unique_identifier(ppi_df)
ppi_df['remap_unique'] = ~ppi_df.duplicated(subset=['xref_partners_sorted'])
print(
    'The act of remapping protein identifiers to UniProt AC, introduced the following duplicates:\n'
)
print(ppi_df.loc[(ppi_df.checker_unique)
                 & ~(ppi_df.remap_unique)].groupby('origin').size(), '\n')

# check if remapping creates new unique values (should never happen)
assert ppi_df.loc[~(ppi_df.original_unique) & (ppi_df.remap_unique), [
    'xref_partners_sorted', 'original_unique', 'remap_unique', 'checker_unique'
]].shape[0] == 0

# report information about duplicates
print('Number of unique interactions per raw dataset:')
print(ppi_df.groupby('origin')['xref_partners_sorted'].nunique(), '\n')
print('Total dataset size:')
print(ppi_df.groupby('origin')['xref_partners_sorted'].size(), '\n')
print('Total number of unique interactions out of {}'.format(ppi_df.shape[0]))
print(np.sum(~ppi_df.duplicated(subset=['xref_partners_sorted'])), '\n')
print(
    'Total number of unique interactions out of {}, where publications are considered unique:'.
    format(ppi_df.shape[0]))
print(
    np.sum(~ppi_df.duplicated(subset=['xref_partners_sorted', 'publication'])),
    '\n')
ppi_df.drop(
    ['original_unique', 'checker_unique', 'remap_unique'],
    axis=1,
    inplace=True)

# remove duplicates
# first sort on dataframe priority
ppi_df['origin'] = pd.Categorical(ppi_df['origin'], [
    'intact-virus-22.03.2018.mitab', 'BIOGRID-ALL-3.4.160.mitab',
    'hpidb2-19.02.2018.mitab', 'virhostnet-01.2018.mitab', 'phi_data.csv'
])
ppi_df.sort_values(by='origin', ascending=True, inplace=True)
size = ppi_df.shape[0]
ppi_df = ppi_df.drop_duplicates(subset=['xref_partners_sorted'], keep='first')
ppi_df = ppi_df.reset_index(drop=True)
print('All duplicate interactions were removed, leaving {} out of {} PPIs.\n'.
      format(ppi_df.shape[0], size))

# remove non-UniProt identifiers
size = ppi_df.shape[0]
ppi_df = ppi_df.loc[(ppi_df.xref_A.str.contains('uniprot'))
                    & (ppi_df.xref_B.str.contains('uniprot'))]
ppi_df = ppi_df.reset_index(drop=True)
print('Omitted {} non-UniProt AC interactions, leaving {} PPIs.'.format(
    size - ppi_df.shape[0], ppi_df.shape[0]))

# Move all host partners in xref_B to xref_A (only an issue for VirHostNet)
# Note, also move ALL other associated labels...
ppi_filter.reorder_pathogen_host_entries(ppi_df, host_taxonids)

# save new ppi dataset
out_ppi = out_path / 'ppi_data/ppi-filter-remap.tsv'
out_ppi.parent.mkdir(parents=True, exist_ok=True)
ppi_df.to_csv(out_ppi, sep='\t', index=False, header=True)
print('\nSaved filtered and remapped PPI dataset to {}'.format(out_ppi))

# save protein list for filtering gaf/interpro files
all_identifiers = pd.Series(
    pd.unique(ppi_df[['xref_A',
                      'xref_B']].values.ravel('K'))).str.split(':').str.get(1)
out_identifiers = out_path / 'ppi_data/uniprot_identifiers.txt'
with out_identifiers.open('w') as out:
    for i in all_identifiers:
        out.write("{}\n".format(i))
print('\nSaved list of all UniProtACs to {}'.format(out_identifiers))
