"""The plotting functions"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['plot_synteny', 'plot_synteny_schematic']

# %% ../nbs/00_core.ipynb 4
import os
from pathlib import Path
from typing import Tuple, Union

import matplotlib.colors as mcolors
import matplotlib.patches as mpatches
import pandas as pd
from matplotlib import pyplot as plt

from . import util

# %% ../nbs/00_core.ipynb 8
def plot_synteny(
    gff,  # a GFF in Pandas dataframe form. Only includes the genes of the syntenic block in question.
    block_start: int = None,  # The start coordinate of the syntenic locus to plot. If None, will be set to the start of the first gene minus 5% of the block length.
    block_end: int = None,  # The end coordinate of the syntenic locus to plot. If None, will be set to the end of the last gene plus 5% of the block length.
    figsize: (float, float) = None,  # the figure size.
    locus_start: str = "start",  # the GFF column that describes the start coordinate of the entities (e.g. gene, mRNA) to be plotted on the chromosome. In a well-behaved GFF, this should be "start".
    locus_end: str = "end",  # the GFF column that describes the end coordinate of the entities (e.g. gene, mRNA) to be plotted on the chromosome. In a well-behaved GFF, this should be "end".
    molecule: str = "seqid",  # the GFF column that holds the molecule name (chromosome/scaffold/contig ID). In a well-behaved GFF, this should be "seqid".
    locus_name: str = "gene_name",  # the GFF column that holds the gene/mRNA symbol. This is usually a tag in the "attributes" column of a well-behaved GFF and should have been extracted prior to this step. It will be plotted according to the `gene_name_offset` parameter.
    fig_title: str = None,  # Figure title.
    save: str | Path = "synteny.svg",  # path to save an SVG image.
    palette: (
        dict | None
    ) = mcolors.CSS4_COLORS,  # color palette in dictionary form. If None, the color column is expected to be a string that can be interpreted as color.
    chromosome_color: str = "lightgray",  # the color of the line representing the molecule the genes are plotted on.
    backup_gene_color: str = "darkgray",  # If the GFF has no column titled "color", the gene arrows will be filled using this color.
    gene_name_offset: float = 0.02,  # Default offset for plotting gene names if the GFF has no column titled "offset".
    chromosome_width: float = 1,  # Linewidth parameter for the line representing the chromosome.
    gene_width: float = 5,  # Width of the arrow patches representing genes. Please refer to the documentation of `matplotlib.pyplot.arrow` for more details.
) -> None:
    "Function to plot a syntenic cluster of genes."

    if len(gff[molecule].unique()) > 1:
        raise ValueError(
            f"There are multiple molecules represented in the input GFF: {gff[molecule].unique()}. Please only input a GFF with genes on a single molecule."
        )

    if block_start is None or block_end is None:
        block_start, block_end = util.syntenic_block_borders(gff)

    if figsize is None:
        figsize = util.estimate_plot_size(gff)

    fig, ax = plt.subplots(figsize=figsize)

    # plot contig
    ax.plot(
        [block_start, block_end],
        [0, 0],
        color=chromosome_color,
        linewidth=chromosome_width,
    )
    for name, gene in gff.iterrows():
        start = gene[locus_start]
        end = gene[locus_end]
        if "color" in gene.index:
            if palette is not None:
                color = mcolors.CSS4_COLORS[gene["color"]]
            else:
                color = gene["color"]
        else:
            color = backup_gene_color

        if "offset" in gene.index:
            offset = gene["offset"]
        else:
            offset = gene_name_offset

        if gene["strand"] == "-":
            arrow_from = end
            arrow_to = start - end - 1000
        else:
            arrow_from = start
            arrow_to = end - start + 1000

        ax.arrow(
            arrow_from,
            0,
            arrow_to,
            0,
            head_width=0.05,
            head_length=1000,
            fc=color,
            ec=color,
            linewidth=gene_width,
        )
        ax.text(
            gene[locus_start] - 10000,
            offset * 20,
            gene[locus_name],
            color=color,
            fontsize=10,
        )

    ax.set_yticks([])
    ax.set_ylim(-1.5, 1)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_visible(False)
    if fig_title is not None:
        ax.set_title(fig_title)
    plt.gca().set_position([0, 0, 1, 1])
    if save is not None:
        if not (isinstance(save, str) | isinstance(save, Path)):
            raise TypeError(
                f"The save location should be a valid Path or string; instead, it is {type(save)}."
            )
        plt.savefig(save, bbox_inches="tight", pad_inches=0)

# %% ../nbs/00_core.ipynb 17
def plot_synteny_schematic(
    gff: pd.DataFrame,  # a GFF in Pandas dataframe form. Only includes the genes of the syntenic block in question. Assumed to be sorted in plotting order.
    block_gene: float = 400,  # length of a gene arrow, in plot coordinate space.
    block_dist: float = 200,  # length of gap between successive genes, in plot coordinate space.
    locus_start: str = "start",  # the GFF column that describes the start coordinate of the entities (e.g. gene, mRNA) to be plotted on the chromosome. In a well-behaved GFF, this should be "start".
    locus_end: str = "end",  # the GFF column that describes the end coordinate of the entities (e.g. gene, mRNA) to be plotted on the chromosome. In a well-behaved GFF, this should be "end".
    molecule: str = "seqid",  # the GFF column that holds the molecule name (chromosome/scaffold/contig ID). In a well-behaved GFF, this should be "seqid".
    locus_name: str = "gene_name",  # the GFF column that holds the gene/mRNA symbol. This is usually a tag in the "attributes" column of a well-behaved GFF and should have been extracted prior to this step. It will be plotted according to the `gene_name_offset` parameter.
    interrupted: bool = False,  # is the syntenic cluster on multiple pseudomolecules?
    palette: (
        dict | None
    ) = mcolors.CSS4_COLORS,  # color palette in dictionary form. If None, the color column is expected to be a string that can be interpreted as color.
    fontsize: int = 30,  # this fontsize works well for default sizes, but is probably too big for smaller plots
    chromosome_color: str = "black",  # the color of the line representing the molecule the genes are plotted on.
    backup_gene_color: str = "darkgray",  # If the GFF has no column titled "color", the gene arrows will be filled using this color.
    gene_name_offset: float = 0.02,  # Default offset for plotting gene names if the GFF has no column titled "offset".
    chromosome_width: float = 3,  # Linewidth parameter for the line representing the chromosome; also, patch linewidth of the arrows representing genes. Please refer to the documentation of `matplotlib.patches.FancyArrowPath` for more details.
    head_length: float = 200,  # Head length for the arrows representing genes. Please refer to the documentation of `matplotlib.patches.FancyArrowPath` for more details.
    head_width: float = 1.7,  # Head width for the arrows representing genes. Please refer to the documentation of `matplotlib.patches.FancyArrowPath` for more details.
    gene_width: float = 1,  # Width of the arrow patches representing genes. Please refer to the documentation of `matplotlib.patches.FancyArrowPath` for more details.
    ax: plt.Axes = None,  # a pre-existing matplotlib axes object, if you wish to include this plot in a composite.
    figsize: (
        float,
        float,
    ) = None,  # the figure size. Will be ignored if ax is not None.
    save: str | Path = "schematic.svg",  # path to save an SVG image.
    return_fig: bool = False,  # whether to return an axes object. Useful for further manipulation of the plot.
) -> plt.Figure | plt.Axes | None:
    # 'Function to plot a syntenic cluster of genes. Each gene is represented by an arrow; left-facing arrows indicate a minus orientation, and right-facing arrows a plus orientation. The genes can be colored using named colors from the `matplotlib.colors.CSS4_COLORS` collection.'

    if len(gff[molecule].unique()) > 1:
        if not interrupted:
            raise ValueError(
                f"There are multiple molecules represented in the input GFF: {gff[molecule].unique()}. Please only input a GFF with genes on a single molecule."
            )

    if figsize is None:
        figsize = util.estimate_plot_size(gff)

    fig, ax = plt.subplots(figsize=figsize)
    # the total length is:
    # - the number of gene blocks
    # - the number of gaps between genes (number of gene blocks - 1)
    # - the start and end gap
    gff_genes = gff[gff[molecule] != ""]
    total_length = (
        (len(gff_genes) * block_gene)
        + (2 * block_dist)
        + ((len(gff_genes) - 1) * block_dist)
    )
    # plot contig
    ax.plot(
        [0, total_length], [0, 0], color=chromosome_color, linewidth=chromosome_width
    )
    start = 0
    end = 0
    genes_seen = 0
    for i, gene in gff.iterrows():
        if gene[molecule] != "":
            genes_seen += 1
        if gene[locus_start] == "":
            # as many gene blocks as we walked, plus as many gaps minus one
            break_start = (genes_seen) * block_gene + (genes_seen + 0.2) * block_dist
            break_end = break_start + 0.6 * block_dist
            ax.plot(
                [break_start, break_end],
                [0, 0],
                color="white",
                linewidth=chromosome_width + 2,
            )
            offset = gene_name_offset
            continue

        start = end + block_dist
        end = start + block_gene
        if "color" in gene.index:
            if gene["color"] == "":
                continue
            if palette is not None:
                color = mcolors.CSS4_COLORS[gene["color"]]
            else:
                color = gene["color"]
        else:
            color = backup_gene_color

        if "offset" in gene.index:
            offset = gene["offset"]
        else:
            offset = gene_name_offset

        # Draws an arrow from (x, y) to (x + dx, y + dy). The width of the arrow is scaled by width
        if gene["strand"] == "-":
            arrow_from = end
            arrow_to = start - end + head_length
        else:
            arrow_from = start
            arrow_to = end - start - head_length

        ec = "black"
        if "edge_color" in gene.index:
            ec = gene["edge_color"]

        arrow = mpatches.FancyArrow(
            arrow_from,
            0,
            arrow_to,
            0,
            width=gene_width,
            head_length=head_length,
            head_width=head_width,
            ec=ec,
            fc=color,
            lw=chromosome_width,
            zorder=2,
            head_starts_at_zero=False,
        )

        if gene["strand"] == ".":
            arrow_from = end
            arrow_to = start - end
            arrow = mpatches.FancyArrow(
                arrow_from,
                0,
                arrow_to,
                0,
                width=gene_width,
                head_length=0,
                head_width=0,
                ec=ec,
                fc=color,
                lw=chromosome_width,
                zorder=2,
                head_starts_at_zero=True,
            )

        ax.add_patch(arrow)
        # ax.arrow(end, 0, start-end-100, 0, head_width=0.05, head_length=100, fc=color, ec=color, linewidth=gene_width)
        ax.text(
            end - block_gene,
            offset * 60,
            gene[locus_name],
            color="black",
            fontsize=fontsize,
        )

    ax.set_yticks([])
    ax.set_xticks([])
    ax.set_ylim(-1.5, 1)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_visible(False)
    ax.spines["bottom"].set_visible(False)

    plt.gca().set_position([0, 0, 1, 1])
    if save is not None:
        if not (isinstance(save, str) | isinstance(save, Path)):
            raise TypeError(
                f"The save location should be a valid Path or string; instead, it is {type(save)}."
            )
        plt.savefig(save, bbox_inches="tight", pad_inches=0)

    if return_fig:
        return ax
