import click
import pandas as pd
from matplotlib import pyplot as plt

from ...modules import plot
from ..parameters import arg_data, arg_output, EWAS_RESULT, INPUT_FILE


@click.group(name="plot")
def plot_cli():
    pass


@plot_cli.command(help="Create a histogram plot of a variable")
@arg_data
@arg_output
@click.argument("Variable", type=click.STRING)
def histogram(data, output, variable):
    # Plot
    plot.histogram(data=data.df, column=variable)
    # Save and Close
    plt.savefig(output)
    plt.close()
    # Log
    click.echo(click.style(f"Done: Saved plot to {output}", fg="green"))


@plot_cli.command(help="Generate a pdf containing distribution plots for each variable")
@arg_data
@arg_output
@click.option(
    "--kind",
    "-k",
    default="count",
    type=click.Choice(["count", "box", "violin", "qq"]),
    help="Kind of plot used for continuous data.  Non-continuous always shows a count plot.",
)
@click.option(
    "--nrows",
    default=4,
    type=click.IntRange(min=1, max=10),
    help="Number of rows per page",
)
@click.option(
    "--ncols",
    default=3,
    type=click.IntRange(min=1, max=10),
    help="Number of columns per page",
)
@click.option(
    "--quality",
    "-q",
    default="medium",
    type=click.Choice(["low", "medium", "high"]),
    help="Quality of the generated plots: low (150 dpi), medium (300 dpi), or high (1200 dpi).",
)
@click.option("--sort/--no-sort", help="Sort variables alphabetically")
def distributions(data, output, kind, nrows, ncols, quality, sort):
    # Plot and save
    plot.distributions(
        data=data.df,
        filename=output,
        continuous_kind=kind,
        nrows=nrows,
        ncols=ncols,
        quality=quality,
        sort=sort,
    )
    # Log
    click.echo(click.style(f"Done: Saved plot to {output}", fg="green"))


@plot_cli.command(help="Generate a manhattan plot of EWAS results")
@click.argument("ewas_result", type=EWAS_RESULT)
@arg_output
@click.option(
    "--categories",
    "-c",
    type=INPUT_FILE,
    default=None,
    help="tab-separate file with two columns: 'Variable' and 'category'",
)
@click.option(
    "--bonferroni",
    type=click.FLOAT,
    default=0.05,
    help="cutoff value to plot bonferroni-adjusted pvalue line",
)
@click.option(
    "--fdr",
    type=click.FLOAT,
    default=None,
    help="cutoff value to plot fdr-adjusted pvalue line",
)
@click.option(
    "--other",
    "-o",
    multiple=True,
    type=EWAS_RESULT,
    help="other datasets to include in the plot",
)
@click.option(
    "--nlabeled",
    default=3,
    type=click.IntRange(min=0, max=50),
    help="label top n points",
)
@click.option(
    "--label",
    default=None,
    multiple=True,
    type=click.STRING,
    help="label points by name",
)
def manhattan(ewas_result, output, categories, bonferroni, fdr, other, nlabeled, label):
    # Load data
    name, data = ewas_result
    data_dict = {name: data}
    for (name, data) in other:
        data_dict[name] = data
    # Load categories, if any
    if categories is not None:
        categories = pd.read_csv(categories, sep="\t")
        categories.columns = ["Variable", "category"]
        categories = categories.set_index("Variable")["category"].to_dict()
    # Plot and save
    plot.manhattan(
        data_dict,
        categories=categories,
        bonferroni=bonferroni,
        fdr=fdr,
        num_labeled=nlabeled,
        label_vars=label,
        filename=output,
    )
    # Log
    click.echo(click.style(f"Done: Saved plot to {output}", fg="green"))


@plot_cli.command(
    help="Generate a manhattan plot of EWAS results showing Bonferroni-corrected pvalues"
)
@click.argument("ewas_result", type=EWAS_RESULT)
@arg_output
@click.option(
    "--categories",
    "-c",
    type=INPUT_FILE,
    default=None,
    help="tab-separate file with two columns: 'Variable' and 'category'",
)
@click.option(
    "--cutoff",
    type=click.FLOAT,
    default=0.05,
    help="cutoff value for plotting the significance line",
)
@click.option(
    "--fdr",
    type=click.FLOAT,
    default=None,
    help="cutoff value to plot Bonferroni-adjusted pvalue line",
)
@click.option(
    "--other",
    "-o",
    multiple=True,
    type=EWAS_RESULT,
    help="other datasets to include in the plot",
)
@click.option(
    "--nlabeled",
    default=3,
    type=click.IntRange(min=0, max=50),
    help="label top n points",
)
@click.option(
    "--label",
    default=None,
    multiple=True,
    type=click.STRING,
    help="label points by name",
)
def manhattan_bonferroni(
    ewas_result, output, categories, cutoff, other, nlabeled, label
):
    # Load data
    name, data = ewas_result
    data_dict = {name: data}
    for (name, data) in other:
        data_dict[name] = data
    # Load categories, if any
    if categories is not None:
        categories = pd.read_csv(categories, sep="\t")
        categories.columns = ["Variable", "category"]
        categories = categories.set_index("Variable")["category"].to_dict()
    # Plot and save
    plot.manhattan_bonferroni(
        data_dict,
        categories=categories,
        cutoff=cutoff,
        num_labeled=nlabeled,
        label_vars=label,
        filename=output,
    )
    # Log
    click.echo(click.style(f"Done: Saved plot to {output}", fg="green"))


@plot_cli.command(
    help="Generate a manhattan plot of EWAS results showing FDR-corrected pvalues"
)
@click.argument("ewas_result", type=EWAS_RESULT)
@arg_output
@click.option(
    "--categories",
    "-c",
    type=INPUT_FILE,
    default=None,
    help="tab-separate file with two columns: 'Variable' and 'category'",
)
@click.option(
    "--cutoff",
    type=click.FLOAT,
    default=0.05,
    help="cutoff value for plotting the significance line",
)
@click.option(
    "--fdr",
    type=click.FLOAT,
    default=None,
    help="cutoff value to plot fdr-adjusted pvalue line",
)
@click.option(
    "--other",
    "-o",
    multiple=True,
    type=EWAS_RESULT,
    help="other datasets to include in the plot",
)
@click.option(
    "--nlabeled",
    default=3,
    type=click.IntRange(min=0, max=50),
    help="label top n points",
)
@click.option(
    "--label",
    default=None,
    multiple=True,
    type=click.STRING,
    help="label points by name",
)
def manhattan_fdr(ewas_result, output, categories, cutoff, other, nlabeled, label):
    # Load data
    name, data = ewas_result
    data_dict = {name: data}
    for (name, data) in other:
        data_dict[name] = data
    # Load categories, if any
    if categories is not None:
        categories = pd.read_csv(categories, sep="\t")
        categories.columns = ["Variable", "category"]
        categories = categories.set_index("Variable")["category"].to_dict()
    # Plot and save
    plot.manhattan_fdr(
        data_dict,
        categories=categories,
        cutoff=cutoff,
        num_labeled=nlabeled,
        label_vars=label,
        filename=output,
    )
    # Log
    click.echo(click.style(f"Done: Saved plot to {output}", fg="green"))


@click.argument("ewas_result", type=EWAS_RESULT)
@arg_output
@click.option(
    "--pvalue_name",
    type=click.Choice(["pvalue", "pvalue_bonferroni", "pvalue_FDR"]),
    default="pvalue",
    help="Which pvalues to use in the plot",
)
@click.option(
    "--cutoff",
    type=click.FLOAT,
    default=0.05,
    help="cutoff value for plotting the significance line",
)
@click.option(
    "--num_rows", type=click.INT, default=20, help="How many EWAS result rows to plot"
)
def top_results(ewas_result, output, pvalue_name, cutoff, num_rows):
    plot.top_results(
        ewas_result,
        pvalue_name=pvalue_name,
        cutoff=cutoff,
        num_rows=num_rows,
        filename=output,
    )
    # Log
    click.echo(click.style(f"Done: Saved plot to {output}", fg="green"))
