"""A parser for Comet output"""
import re
import logging

import pandas as pd

from .txt import read_txt
from .. import utils

LOGGER = logging.getLogger(__name__)


def read_comet(
    txt_files, pairing_file_name=None, decoy_prefix="DECOY_", copy_data=True
):
    """Read peptide-spectrum matches (PSMs) from Comet output.
    Can parse tab-delimited files.

    Parameters
    ----------
    txt_files : str, pandas.DataFrame or tuple of str
        One or more collection of PSMs in the Crux tab-delimited format.
    pairing_file_name : str, optional
        A tab-delimited file that explicity pairs target and decoy peptide
        sequences. Requires one column labled 'target' that contains target
        sequences and a second colun labeled 'decoy' that contains decoy
        sequences. This file can be generated by setting --peptide-list=T
        in tide-index.
    decoy_prefix : str, optional
        The prefix used to indicate a decoy protein in the protein column.
        Default value is 'DECOY_'.
    copy_data : bool, optional
        If true, a deep copy of the data is created. This uses more memory, but
        is safer because it prevents accidental modification of the underlying
        data. This argument only has an effect when `txt_files` is a
        :py:class:`pandas.DataFrame`

    Returns
    -------
    PsmDataset
        A :py:class:`~crema.dataset.PsmDataset` object containing the parsed
        PSMs.
    """
    target = "target/decoy"
    peptide = "modified_peptide"
    spectrum = ["scan", "exp_neutral_mass"]  # No file col to use
    pairing = ""
    protein = "protein"
    protein_delim = ","

    # Possible score columns output by Comet.
    scores = {
        "e-value",
        "xcorr",
        "sp_score",
    }
    scores_all = scores

    # Keep only Comet scores that exist in all of the files.
    skip_first_line = False
    crux_comet = False  # tracks whether standalone comet or Crux Comet
    if isinstance(txt_files, pd.DataFrame):
        scores = scores.intersection(set(txt_files.columns))
    else:
        txt_files = utils.listify(txt_files)
        for txt_file in txt_files:
            with open(txt_file) as txt_ref:
                # First line of Comet output consists only of version
                # If statement below in case first line is removed
                line = txt_ref.readline().rstrip()
                if line.startswith("CometVersion"):
                    line = txt_ref.readline().rstrip()
                    skip_first_line = True
                cols = line.split("\t")
                scores = scores.intersection(set(cols))

                # Standalone Comet and Crux Comet use different column headers
                if protein not in cols:
                    protein = "protein id"
                    peptide = "modified sequence"
                    spectrum = ["scan", "spectrum neutral mass"]
                    crux_comet = True

    if not scores:
        raise ValueError(
            "Could not find any of the Comet score columns in all of the files."
            f" The columns Crema looks for are {', '.join(list(scores_all))}"
        )

    scores = list(scores)

    # Read in the files:
    fields = [*spectrum, peptide, target, *scores, pairing, protein]
    if isinstance(txt_files, pd.DataFrame):
        data = txt_files.copy(deep=copy_data).loc[:, fields]
    else:
        data = pd.concat(
            [
                utils.parse_psms_txt(f, fields, skip_first_line)
                for f in txt_files
            ]
        )

    if crux_comet and decoy_prefix == "DECOY_":
        decoy_prefix = "decoy_"
    data["target/decoy"] = ~data[protein].str.contains(decoy_prefix)

    psms = read_txt(
        data,
        target_column=target,
        spectrum_columns=spectrum,
        score_columns=scores,
        peptide_column=peptide,
        protein_column=protein,
        protein_delim=protein_delim,
        sep="\t",
        copy_data=False,
    )

    # Remove first and last amino acid from sequence
    # Looks like "R.WVNEK.Y"
    peptide_column = psms.peptides
    new_peptide_column = peptide_column.str[2:-2]
    psms.set_peptide_column(new_peptide_column)

    # always pair target and decoys for Comet
    if pairing_file_name == None:
        # implicit pairing based off fact that Comet reverses peptides
        psms._peptide_pairing = _create_pairing(
            data, peptide, protein, decoy_prefix
        )
    else:  # explicit pairing
        psms._peptide_pairing = utils.create_pairing_from_file(
            pairing_file_name
        )

    # Remove decoy prefix from protein ID
    protein_column = psms.proteins
    new_protein_column = protein_column.str.replace(
        decoy_prefix, "", regex=True
    )
    psms.set_protein_column(new_protein_column)

    return psms


def _create_pairing(pairing_data, peptide_col, protein_col, decoy_prefix):
    """Parse a single Comet dataframe to implicity pair target and
    decoy sequences.

    Parameters
    ----------
    pairing_data : pandas.DataFrame
        A collection of PSMs with the necessary columns to create a
        target/decoy peptide pairing. Required columns are modified sequence
        and protein ID.
    peptide_col: str
        Identifies the modified sequence column.
    protein_col: str
        Identifies the protein ID column.
    decoy_prefix: str
        The prefix used to indicate a decoy protein in the protein column.

    Returns
    -------
    pairing : dict
        A map of target and decoy peptide sequence pairings. Targets with
        missing decoys will not be included among the keys.

    """
    # ensure pairing_data dataframe contains all necessary columns
    req_fields = [peptide_col, protein_col]

    if not set(req_fields).issubset(pairing_data.columns):
        miss = ", ".join(set(req_fields) - set(pairing_data.columns))
        raise ValueError(
            f"Required columns for peptide pairing were not detected: {miss}"
        )

    pairing_data = pairing_data.loc[:, req_fields]
    pairing_data[peptide_col] = pairing_data[peptide_col].str[2:-2]

    pairing_data["target/decoy"] = ~pairing_data[protein_col].str.contains(
        decoy_prefix
    )

    reverse_peptide_list = []
    for seq in list(pairing_data[peptide_col]):
        seq_sp = re.split(r"(?<=.)(?=[A-Z])", seq)
        peptide_rev = "".join([seq_sp[0], *reversed(seq_sp[1:-1]), seq_sp[-1]])
        reverse_peptide_list.append(peptide_rev)

    pairing_data["reverse_peptide"] = reverse_peptide_list

    targets = pairing_data[pairing_data["target/decoy"]]
    decoys = pairing_data[~pairing_data["target/decoy"]]

    dic1 = dict(zip(targets[peptide_col], targets["reverse_peptide"]))
    dic2 = dict(zip(decoys["reverse_peptide"], decoys[peptide_col]))

    dic2.update(dic1)
    return dic2
