from logging import error, warning

from matplotlib.axes import Axes
from matplotlib.figure import Figure, SubFigure
import rdkit.Chem.Draw as rdDraw
import rdkit
import rdkit.Chem as rdChem

import matplotlib as mpl

from shnitsel.analyze.pca import PCAResult
from shnitsel.bridges import to_mol
from shnitsel.data.dataset_containers import Frames

from .common import centertext, label_plot_grid
from ...plot.common import figax, mpl_imshow_png, mpl_svg_into_axes

import xarray as xr


def mol_to_png(mol: rdChem.Mol, width: int = 320, height: int = 240) -> bytes:
    """Helper function to plot an rdkit.Mol object as a 2D structural representation

    Parameters
    ----------
    mol : rdChem.Mol
            RDKit representation of the molecule to plot
    width : int, optional
            Width of the canvas to plot to. Defaults to 320.
    height : int, optional
            Height of the canvas to plot to. Defaults to 240.

    Returns
    -------
    bytes
        The drawn structure of the `mol` molecule in 2D as a sequence of bytes.
    """

    d = rdDraw.rdMolDraw2D.MolDraw2DCairo(width, height)

    d.drawOptions().setBackgroundColour((1, 1, 1, 0))
    d.drawOptions().padding = 0.05

    d.DrawMolecule(mol)
    d.FinishDrawing()
    return d.GetDrawingText()


def mol_to_svg(mol: rdChem.Mol, width: int = 320, height: int = 240) -> str:
    """Helper function to plot an rdkit.Mol object as a 2D structural representation using vector graphics

    Parameters
    ----------
    mol : rdChem.Mol
            RDKit representation of the molecule to plot
    width : int, optional
            Width of the canvas to plot to. Defaults to 320.
    height : int, optional
            Height of the canvas to plot to. Defaults to 240.

    Returns
    -------
    bytes
        The drawn structure of the `mol` molecule in 2D as a sequence of bytes.
    """

    d = rdDraw.rdMolDraw2D.MolDraw2DSVG(width, height)

    d.drawOptions().setBackgroundColour((1, 1, 1, 0))
    d.drawOptions().padding = 0.05

    d.DrawMolecule(mol)
    d.FinishDrawing()
    return d.GetDrawingText()


def format_inchi(inchi: str, max_line_length: int = 30) -> str:
    """Helper function to split an inchi string across
    multiple lines to not overlap other plots in the DataSheet pages.

    Parameters
    ----------
    inchi : str
            InChi string as generated by RDKit
    max_line_length : int, optional
            Maximum line length. Defaults to 30 characters

    Returns
    -------
    str
        The string broken into multiple lines with inserted \\n characters.
    """
    if len(inchi) < max_line_length:
        return inchi
    else:
        split = inchi.split('/')
        if len(split) not in {4, 5}:
            warning(f"Unexpected InChi: {split=}")
        lens = [len(s) for s in split]
        split[2] = '\n' + split[2]
        if sum(lens[2:]) > max_line_length:
            split[3] = '\n' + split[3]
        return '/'.join(split)


def plot_structure(
    mol: rdChem.Mol,
    name: str = '',
    smiles: str | None = None,
    inchi: str | None = None,
    fig: Figure | SubFigure | None = None,
    ax: Axes | None = None,
) -> Axes:
    """Plot function to create a 2d structure representation of a molecule `mol` using rdkit and create a graph with
    a 2d image representation into `fig` or `ax` with at least one of the two needing to be provided.

    Parameters
    ----------
    mol : rdChem.Mol
            The RDKit molecule object to visualize.
    name : str, optional
            The name of the structure. Would be used as a plot title if not empty. Defaults to ''.
    smiles : str | None, optional
            Smiles strings to include in the x-axis label. Defaults to None.
    inchi : str | None, optional
            InChI strings to include in the x-axis label. Defaults to None.
    fig : Figure | SubFigure | None, optional
            Figure to plot the graphic into if `ax` not provided. Defaults to None.
    ax : Axes | None, optional
            Axis to plot the graphic into. Defaults to None, meaning they will be created from `fig`.

    Returns
    -------
    Axes
        The axes the graphic was plotted to.
    """
    fig, ax = figax(fig, ax)
    try:
        # png = mol_to_png(mol)
        svg = mol_to_svg(mol)
        # print(svg)
    except ImportError as err:
        error(
            "ImportError from rdkit.Chem.Draw while "
            f"attempting to plot structure: {err}"
        )
        centertext("ImportError while attempting\nto plot structure", ax)
        return ax

    # mpl_imshow_png(ax, png)
    mpl_svg_into_axes(ax, svg)

    if len(name) > 0:
        ax.set_title(name)

    ax.axis('on')
    ax.get_yaxis().set_visible(False)
    ax.tick_params(axis="x", bottom=False, labelbottom=False)
    xlabel_string = ""
    # Add smiles to label if provided
    if smiles is not None:
        xlabel_string = f"SMILES: '{smiles}'"

    # Add InChI to label if provided
    if inchi is not None:
        inchi = format_inchi(inchi)
        if len(xlabel_string) > 0:
            xlabel_string += '\n'
        xlabel_string += f"{inchi}"

    ax.set_xlabel(xlabel_string, fontsize='small')
    # axy.tick_params(axis="y", labelleft=False)
    return ax


def plot_pca_structure(
    frames: Frames | xr.Dataset,
    pca_data: xr.DataArray,
    fig: Figure | SubFigure | None = None,
    axs: dict[tuple[int, str], Axes] | None = None,
) -> dict[tuple[int, str], Axes]:
    """Function to plot the most extreme pca combinations as structure plots.

    Parameters
    ----------
    frames : Frames | xr.Dataset
        The raw input data to use for extracting the positional data and structure.
    pca_data : xr.DataArray
        PCA decomposition data.
    fig : Figure | SubFigure, optional
        Figure to create axes into. If `axs` are provided, this can be left empty.
    axs : dict[tuple[int,str],Axes] | None, optional
        Axes to plot the extrema to. One row per pca, one column for 'min' or 'max'. Defaults to None.

    Returns
    -------
    dict[tuple[int,str],Axes]
        The axes the extrema have been plotted to.
    """
    assert 'PC' in pca_data.sizes, f"pca_data argument is no result of PCA: {pca_data}"
    if isinstance(frames, xr.Dataset):
        try:
            frames = Frames(frames)
        except:
            raise ValueError("Provided dataset is no appropriate frameset")

    assert frames.has_variable('atXYZ'), "No positional data provided."

    combined_ds = frames.dataset.assign(pca_data=pca_data)
    num_pca = combined_ds.sizes['PC']
    if axs is None:
        if fig is not None:
            tmp_ax = fig.subplots(num_pca, 2)
            axs = {}
            for i, label in enumerate(['min', 'max']):
                for j in range(num_pca):
                    axs[(j, label)] = tmp_ax[j, i]
        else:
            raise ValueError("Please provide either a figure or axes.")

    pc_labels = []
    col_labels = ['min', 'max']
    for j in range(num_pca):
        single_pc_data = combined_ds.isel(PC=j)
        min_struct = single_pc_data.pca_data.idxmin(skipna=True)
        max_struct = single_pc_data.pca_data.idxmax(skipna=True)

        # print(min_struct, max_struct)s

        pos_argmin = combined_ds.sel({min_struct.name: min_struct.values}).atXYZ
        pos_argmax = combined_ds.sel({max_struct.name: max_struct.values}).atXYZ

        min_ax = axs[(j, 'min')]
        max_ax = axs[(j, 'max')]

        min_mol = to_mol(pos_argmin, to2D=True)
        max_mol = to_mol(pos_argmax, to2D=True)

        plot_structure(min_mol, ax=min_ax)
        plot_structure(max_mol, ax=max_ax)
        pc_labels.append(f"PC{j + 1}")
        min_ax.get_yaxis().set_visible(True)
        min_ax.set_yticks([])

    if fig is not None:
        label_plot_grid(fig, row_headers=pc_labels, col_headers=col_labels)

    return axs
