# -*- coding: utf-8 -*-

"""Command line interface."""

import json
import logging
from collections import defaultdict

import click
import numpy as np
import pandas as pd
from tqdm import tqdm

from pemt.chemical_extractor.experimental_data_extraction import extract_chemicals
from pemt.constants import MAPPER_DIR, PATENT_DIR
from pemt.patent_extractor.patent_chemical_harmonizer import harmonize_chemicals
from pemt.patent_extractor.patent_enrichment import extract_patent

logger = logging.getLogger(__name__)


@click.group()
def main():
    """Run PEMT."""
    logging.basicConfig(format="%(asctime)s - %(levelname)s - %(name)s - %(message)s")


input_data = click.option(
    "--data",
    help="Path to tab-separated gene data file",
    type=click.Path(file_okay=True, dir_okay=False, exists=True),
    required=True,
)
input_data_type = click.option(
    "--input-type",
    help="Type of data file i.e. 'tab' for tsv or 'comma' for csv files",
    type=str,
    default=",",
)
analysis_name = click.option(
    "--name",
    help="Name of the analysis that is to be run",
    type=str,
    required=True,
)
has_uniprot = click.option(
    "--uniprot/--no-uniprot",
    default=True,
    help="Boolean value indicating whether the gene data file has uniprot ids or not.",
)
system_name = click.option(
    "--os",
    type=click.Choice(["linux", "mac", "windows"], case_sensitive=False),
    help="The OS system on which is the script is running",
    multiple=False,
)
chromedriver_path = click.option(
    "--chromedriver-path",
    help="The path where the chromedriver can be found on the users computer",
    type=str,
    required=True,
)
patent_year = click.option(
    "--year",
    help="The year from which you want to retrive patents from",
    type=int,
    default=2000,
)
from_chemical = click.option(
    "--chemical/--no-chemical",
    default=True,
    help="Boolean value indicating whether the chemical data is provided by the user or not",
)
chemcial_data = click.option(
    "--chemical-data",
    help="Path to tab-separated chemical data file with single column of chembl_ids",
    type=click.Path(),
    default="",
)


@main.command(help="Extract chemicals for genes of interest")
@analysis_name
@input_data
@input_data_type
@has_uniprot
def run_chemical_extractor(
    name: str, data: str, input_type: str, uniprot: bool
) -> None:
    """Extracting chemicals for genes with experiemtal data."""
    click.echo(f"Starting the chemical extractor pipeline for {name}")

    if uniprot:
        with_uniprot = True
    else:
        with_uniprot = False

    gene_chemical_dict = extract_chemicals(
        analysis_name=name,
        gene_file_path=data,
        file_separator=input_type,
        is_uniprot=with_uniprot,
    )

    click.echo(
        f"Completed the chemical extractor pipeline for {len(gene_chemical_dict)} genes."
    )
    click.echo(f"Data file can be found under {MAPPER_DIR}")


@main.command(help="Extract patent for filtered chemicals")
@analysis_name
@system_name
@chromedriver_path
@patent_year
@from_chemical
@chemcial_data
def run_patent_extractor(
    name: str,
    os: str,
    chromedriver_path: str,
    year: str,
    chemical: bool,
    chemical_data: str,
) -> None:
    """Extracting patent from chemical data."""
    click.echo(f"Starting to pre-process the chemical data for patent retrieval")

    if chemical:
        df = pd.read_csv(chemical_data, sep="\t", dtype=str)

        df.to_csv(f"{PATENT_DIR}/{name}_chemicals.tsv", sep="\t", index=False)

        harmonize_chemicals(analysis_name=name, from_genes=False)
    else:
        harmonize_chemicals(analysis_name=name)

    click.echo(f"Starting the patent extractor pipeline for {name}")

    patent_df = extract_patent(
        analysis_name=name,
        chrome_driver_path=chromedriver_path,
        os_system=os,
        patent_year=year,
    )

    if patent_df.empty:
        click.echo(f"No patents found!")
        return None

    # Since the original patent data has chemical with no patents, we remove those entries from the data
    patent_df = patent_df[~patent_df["patent_id"].isna()]
    patent_df.to_csv(
        f"{PATENT_DIR}/cleaned_{name}_patent_data.tsv", sep="\t", index=False
    )

    if chemical:
        click.echo(f"Done with retrival of patents")
        click.echo(f"Data file can be found under {PATENT_DIR}")
        return

    gene_chemical_data = json.load(open(f"{MAPPER_DIR}/{name}_gene_to_chemicals.json"))
    chemical_to_gene_mapper = defaultdict(set)

    for gene, chemicals in tqdm(gene_chemical_data.items()):
        for chemical in chemicals:
            chemical_to_gene_mapper[chemical].add(gene)

    patent_df["genes"] = patent_df["chembl"].map(
        lambda x: ", ".join(chemical_to_gene_mapper[x])
    )

    patent_df.to_csv(f"{PATENT_DIR}/{name}_gene_patent_data.tsv", sep="\t", index=False)

    click.echo(f"Done with retrival of patents")
    click.echo(f"Data file can be found under {PATENT_DIR}")


@main.command(help="Run the PEMT tool with gene data")
@analysis_name
@input_data
@input_data_type
@has_uniprot
@chromedriver_path
@system_name
@patent_year
def run_pemt(
    name: str,
    data: str,
    input_type: str,
    uniprot: bool,
    chromedriver_path: str,
    os: str,
    year: str,
) -> None:
    """Runs the PEMT tool with all the components together."""
    click.echo(f"Starting to run PEMT workflow for {name}")

    click.echo(f"Running the chemical extractor pipeline")

    if uniprot:
        with_uniprot = True
    else:
        with_uniprot = False

    gene_chemical_dict = extract_chemicals(
        analysis_name=name,
        gene_file_path=data,
        file_separator=input_type,
        is_uniprot=with_uniprot,
    )

    click.echo(
        f"Completed running the chemical extractor pipeline for {len(gene_chemical_dict)} genes."
    )

    click.echo(f"Ppre-processing the chemical data for patent retrieval")

    harmonize_chemicals(analysis_name=name)

    click.echo(f"Running the patent extractor pipeline")

    patent_df = extract_patent(
        analysis_name=name,
        chrome_driver_path=chromedriver_path,
        os_system=os,
        patent_year=year,
    )

    if patent_df.empty:
        click.echo(f"No patents found!")
        return None

    # Since the original patent data has chemical with no patents, we remove those entries from the data
    patent_df = patent_df[~patent_df["patent_id"].isna()]
    patent_df.to_csv(
        f"{PATENT_DIR}/cleaned_{name}_patent_data.tsv", sep="\t", index=False
    )

    chemical_to_gene_mapper = defaultdict(set)

    for gene, chemicals in tqdm(gene_chemical_data.items()):
        for chemical in chemicals:
            chemical_to_gene_mapper[chemical].add(gene)

    patent_df["genes"] = patent_df["chembl"].map(
        lambda x: ", ".join(chemical_to_gene_mapper[x])
    )

    patent_df.to_csv(f"{PATENT_DIR}/{name}_gene_patent_data.tsv", sep="\t", index=False)

    click.echo(f"Done with retrival of patents")


if __name__ == "__main__":
    main()
