from collections import OrderedDict
import os.path as op
import logging

import numpy as np
import imageio as io
from PIL import Image, ImageChops

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.transforms as mtransforms

import nibabel as nib
import dipy.tracking.streamlinespeed as dps
from dipy.align import resample

import AFQ.utils.volume as auv
import AFQ.registration as reg
import AFQ.utils.streamlines as aus

__all__ = ["Viz"]

viz_logger = logging.getLogger("AFQ")
tableau_20 = [
    (0.12156862745098039, 0.4666666666666667, 0.7058823529411765),
    (0.6823529411764706, 0.7803921568627451, 0.9098039215686274),
    (1.0, 0.4980392156862745, 0.054901960784313725),
    (1.0, 0.7333333333333333, 0.47058823529411764),
    (0.17254901960784313, 0.6274509803921569, 0.17254901960784313),
    (0.596078431372549, 0.8745098039215686, 0.5411764705882353),
    (0.8392156862745098, 0.15294117647058825, 0.1568627450980392),
    (1.0, 0.596078431372549, 0.5882352941176471),
    (0.5803921568627451, 0.403921568627451, 0.7411764705882353),
    (0.7725490196078432, 0.6901960784313725, 0.8352941176470589),
    (0.5490196078431373, 0.33725490196078434, 0.29411764705882354),
    (0.7686274509803922, 0.611764705882353, 0.5803921568627451),
    (0.8901960784313725, 0.4666666666666667, 0.7607843137254902),
    (0.9686274509803922, 0.7137254901960784, 0.8235294117647058),
    (0.4980392156862745, 0.4980392156862745, 0.4980392156862745),
    (0.7803921568627451, 0.7803921568627451, 0.7803921568627451),
    (0.7372549019607844, 0.7411764705882353, 0.13333333333333333),
    (0.8588235294117647, 0.8588235294117647, 0.5529411764705883),
    (0.09019607843137255, 0.7450980392156863, 0.8117647058823529),
    (0.6196078431372549, 0.8549019607843137, 0.8980392156862745)]
tableau_extension = [
    (0.95, 0.85, 0.25),  # Warm, Soft Yellow
    (0.98, 0.92, 0.5),
    (0.4, 0.5, 0.6),  # Cool, Muted Blue
    (0.6, 0.7, 0.9)]
large_font = 28
medium_font = 24
small_font = 20
marker_size = 200

COLOR_DICT = OrderedDict({
    "Left Anterior Thalamic": tableau_20[0], "C_L": tableau_20[0],
    "Right Anterior Thalamic": tableau_20[1], "C_R": tableau_20[1],
    "Left Corticospinal": tableau_20[2],
    "Right Corticospinal": tableau_20[3],
    "Left Cingulum Cingulate": tableau_20[4], "MCP": tableau_20[4],
    "Right Cingulum Cingulate": tableau_20[5], "CCMid": tableau_20[5],
    "Forceps Minor": tableau_20[8], "CC_ForcepsMinor": tableau_20[8],
    "Forceps Major": tableau_20[9], "CC_ForcepsMajor": tableau_20[9],
    "Left Inferior Fronto-occipital": tableau_20[10],
    "IFOF_L": tableau_20[10],
    "Right Inferior Fronto-occipital": tableau_20[11],
    "IFOF_R": tableau_20[11],
    "Left Inferior Longitudinal": tableau_20[12], "F_L": tableau_20[12],
    "Right Inferior Longitudinal": tableau_20[13], "F_R": tableau_20[13],
    "Left Superior Longitudinal": tableau_20[14],
    "Right Superior Longitudinal": tableau_20[15],
    "Left Uncinate": tableau_20[16], "UF_L": tableau_20[16],
    "Right Uncinate": tableau_20[17], "UF_R": tableau_20[17],
    "Left Arcuate": tableau_20[18], "AF_L": tableau_20[18],
    "Right Arcuate": tableau_20[19], "AF_R": tableau_20[19],
    "Left Posterior Arcuate": tableau_20[6],
    "Right Posterior Arcuate": tableau_20[7],
    "Left Vertical Occipital": tableau_extension[0],
    "Right Vertical Occipital": tableau_extension[1],
    "median": tableau_20[6],
    # Paul Tol's palette for callosal bundles
    "Callosum Orbital": (0.2, 0.13, 0.53),
    "Callosum Anterior Frontal": (0.07, 0.47, 0.2),
    "Callosum Superior Frontal": (0.27, 0.67, 0.6),
    "Callosum Motor": (0.53, 0.8, 0.93),
    "Callosum Superior Parietal": (0.87, 0.8, 0.47),
    "Callosum Posterior Parietal": (0.8, 0.4, 0.47),
    "Callosum Occipital": (0.67, 0.27, 0.6),
    "Callosum Temporal": (0.53, 0.13, 0.33)})

POSITIONS = OrderedDict({
    "Left Anterior Thalamic": (1, 0), "Right Anterior Thalamic": (1, 4),
    "C_L": (1, 0), "C_R": (1, 4),
    "Left Corticospinal": (1, 1), "Right Corticospinal": (1, 3),
    "Left Cingulum Cingulate": (3, 1),
    "Right Cingulum Cingulate": (3, 3),
    "MCP": (3, 1), "CCMid": (3, 3),
    "Forceps Minor": (4, 2), "Forceps Major": (0, 2),
    "CC_ForcepsMinor": (4, 2), "CC_ForcepsMajor": (0, 2),
    "Left Inferior Fronto-occipital": (4, 1),
    "Right Inferior Fronto-occipital": (4, 3),
    "IFOF_L": (4, 1), "IFOF_R": (4, 3),
    "Left Inferior Longitudinal": (3, 0),
    "Right Inferior Longitudinal": (3, 4),
    "F_L": (3, 0), "F_R": (3, 4),
    "Left Superior Longitudinal": (2, 1),
    "Right Superior Longitudinal": (2, 3),
    "Left Arcuate": (2, 0), "Right Arcuate": (2, 4),
    "AF_L": (2, 0), "AF_R": (2, 4),
    "Left Uncinate": (0, 1), "Right Uncinate": (0, 3),
    "UF_L": (0, 1), "UF_R": (0, 3)})

CSV_MAT_2_PYTHON = \
    {'fa': 'dti_fa', 'md': 'dti_md'}

SCALE_MAT_2_PYTHON = \
    {'dti_md': 0.001}

SCALAR_REMOVE_MODEL = \
    {'dti_md': 'MD', 'dki_md': 'MD', 'dki_fa': 'FA', 'dti_fa': 'FA'}

RECO_FLIP = ["IFO_L", "IFO_R", "UNC_L", "ILF_L", "ILF_R"]

BEST_BUNDLE_ORIENTATIONS = {
    "Left Anterior Thalamic": ("Sagittal", "Left"),
    "Right Anterior Thalamic": ("Sagittal", "Right"),
    "Left Corticospinal": ("Sagittal", "Left"),
    "Right Corticospinal": ("Sagittal", "Right"),
    "Left Cingulum Cingulate": ("Sagittal", "Left"),
    "Right Cingulum Cingulate": ("Sagittal", "Right"),
    "Forceps Minor": ("Axial", "Top"),
    "Forceps Major": ("Axial", "Top"),
    "Left Inferior Fronto-occipital": ("Sagittal", "Left"),
    "Right Inferior Fronto-occipital": ("Sagittal", "Right"),
    "Left Inferior Longitudinal": ("Sagittal", "Left"),
    "Right Inferior Longitudinal": ("Sagittal", "Right"),
    "Left Superior Longitudinal": ("Axial", "Top"),
    "Right Superior Longitudinal": ("Axial", "Top"),
    "Left Uncinate": ("Axial", "Bottom"),
    "Right Uncinate": ("Axial", "Bottom"),
    "Left Arcuate": ("Sagittal", "Left"),
    "Right Arcuate": ("Sagittal", "Right"),
    "Left Vertical Occipital": ("Coronal", "Back"),
    "Right Vertical Occipital": ("Coronal", "Back"),
    "Left Posterior Arcuate": ("Coronal", "Back"),
    "Right Posterior Arcuate": ("Coronal", "Back")}


class PanelFigure():
    """
    Super useful class for organizing existing images
    into subplots using matplotlib
    """

    def __init__(self, num_rows, num_cols, width, height,
                 panel_label_kwargs={}):
        """
        Initialize PanelFigure.

        Parameters
        ----------
        num_rows : int
            Number of rows in figure
        num_cols : int
            Number of columns in figure
        width : int
            Width of figure in inches
        height : int
            Height of figure in inches
        panel_label_kwargs : dict
            Additional arguments for matplotlib's text method,
            which is used to add panel labels to each subplot
        """
        self.fig = plt.figure(figsize=(width, height))
        self.grid = plt.GridSpec(num_rows, num_cols, hspace=0, wspace=0)
        self.subplot_count = 0
        self.panel_label_kwargs = dict(
            fontfamily="Helvetica-Bold",
            fontsize="xx-large",
            color="white",
            fontweight='bold',
            verticalalignment="top",
            bbox=dict(
                facecolor='none',
                edgecolor='none'))
        self.panel_label_kwargs.update(panel_label_kwargs)

    def add_img(self, fname, x_coord, y_coord, reduct_count=1,
                subplot_label_pos=(0.1, 1.0), legend=None, legend_kwargs={},
                add_panel_label=True, panel_label_kwargs={}):
        """
        Add image from fname into figure as a panel.

        Parameters
        ----------
        fname : str
            path to image file to add to subplot
        x_coord : int or slice
            x coordinate(s) of subplot in matlotlib figure
        y_coord : int or slice
            y coordinate(s) of subplot in matlotlib figure
        reduct_count : int
            number of times to trim whitespace around image
            Default: 1
        subplot_label_pos : tuple of floats
            position of subplot label
            Default: (0.1, 1.0)
        legend : dict
            dictionary of legend items, where keys are labels
            and values are colors
            Default: None
        legend_kwargs : dict
            ADditional arguments for matplotlib's legend method
        add_panel_label : bool
            Whether or not to add a panel label to the subplot
            Default: True
        panel_label_kwargs : dict
            Additional arguments for matplotlib's text method,
            which is used to add panel labels to each subplot.
        """
        ax = self.fig.add_subplot(self.grid[y_coord, x_coord])
        im1 = Image.open(fname)
        for _ in range(reduct_count):
            im1 = trim(im1)
        if legend is not None:
            patches = []
            for value, color in legend.items():
                patches.append(mpatches.Patch(
                    color=color,
                    label=value))
            ax.legend(handles=patches, borderaxespad=0., **legend_kwargs)
        if add_panel_label:
            trans = mtransforms.ScaledTranslation(
                10 / 72, -5 / 72, self.fig.dpi_scale_trans)
            this_pl_kwargs = self.panel_label_kwargs.copy()
            this_pl_kwargs.update(panel_label_kwargs)
            ax.text(
                subplot_label_pos[0], subplot_label_pos[1],
                f"{chr(65+self.subplot_count)}",
                transform=ax.transAxes + trans,
                **this_pl_kwargs)
        ax.imshow(np.asarray(im1), aspect=1)
        ax.axis('off')
        self.subplot_count = self.subplot_count + 1
        return ax

    def format_and_save_figure(self, fname, trim_final=True):
        """
        Format and save figure to fname.
        Parameters
        ----------
        fname : str
            Path to save figure to
        trim : bool
            Whether or not to trim whitespace around figure.
            Default: True
        """
        self.fig.tight_layout()
        self.fig.savefig(fname, dpi=300)
        if trim_final:
            im1 = Image.open(fname)
            im1 = trim(im1)
            im1.save(fname, dpi=(300, 300))


def get_eye(view, direc):
    direc = direc.lower()
    view = view.lower()

    if view in ["sagital", "saggital"]:
        viz_logger.warning("You don't know how to spell sagggitttal!")
        view = "sagittal"

    if view not in ["sagittal", "coronal", "axial"]:
        raise ValueError(
            "View must be one of: sagittal, coronal, or axial")

    if direc not in ["left", "right", "top", "bottom", "front", "back"]:
        raise ValueError(
            "View must be one of: left, right, top, bottom, front, back")

    eye = {}
    if view == "sagittal":
        if direc == "left":
            eye["x"] = -1
        else:
            eye["x"] = 1
        eye["y"] = 0
        eye["z"] = 0
    elif view == "coronal":
        eye["x"] = 0
        if direc == "front":
            eye["y"] = 1
        else:
            eye["y"] = -1
        eye["z"] = 0
    elif view == "axial":
        eye["x"] = 0
        eye["y"] = 0
        if direc == "top":
            eye["z"] = 1
        else:
            eye["z"] = -1
    return eye


def display_string(scalar_name):
    if isinstance(scalar_name, str):
        return scalar_name.replace("_", " ").upper()
    else:
        return [sn.replace("_", " ").upper() for sn in scalar_name]


def gen_color_dict(bundles):
    """
    Helper function.
    Generate a color dict given a list of bundles.
    """
    def incr_color_idx(color_idx):
        return (color_idx + 1) % 20
    custom_color_dict = {}
    color_idx = 0
    for bundle in bundles:
        if bundle not in custom_color_dict.keys():
            if bundle in COLOR_DICT.keys():
                custom_color_dict[bundle] = COLOR_DICT[bundle]
            else:
                other_bundle = bundle
                if bundle.startswith("Left "):
                    other_bundle = "Right" + other_bundle[5:]
                elif bundle.startswith("Right "):
                    other_bundle = "Left" + other_bundle[4:]
                other_bundle = str(other_bundle)

                if other_bundle == bundle:  # lone bundle
                    custom_color_dict[bundle] = tableau_20[color_idx]
                    color_idx = incr_color_idx(color_idx)
                else:  # right left pair
                    if color_idx % 2 != 0:
                        color_idx = incr_color_idx(color_idx)
                    custom_color_dict[bundle] =\
                        tableau_20[color_idx]
                    custom_color_dict[other_bundle] =\
                        tableau_20[color_idx + 1]
                    color_idx = incr_color_idx(incr_color_idx(color_idx))
    return custom_color_dict


def viz_import_msg_error(module):
    """Alerts user to install the appropriate viz module """
    if module == "plot":
        msg = "To make plots in pyAFQ, you will need to install "
        msg += "the relevant plotting software packages."
        msg += "You can do that by installing pyAFQ with "
        msg += "`pip install AFQ[plot]`"
    else:
        msg = f"To use {module.upper()} visualizations in pyAFQ, you will "
        msg += f"need to have {module.upper()} installed. "
        msg += "You can do that by installing pyAFQ with "
        msg += f"`pip install AFQ[{module.lower()}]`, or by "
        msg += f"separately installing {module.upper()}: "
        msg += f"`pip install {module.lower()}`."
    return msg


def tract_generator(trk_file, bundle, colors, n_points,
                    n_sls_viz=3600, n_sls_min=75):
    """
    Generates bundles of streamlines from the tractogram.
    Only generates from relevant bundle if bundle is set.
    Otherwise, returns all streamlines.

    Helper function

    Parameters
    ----------
    trk_file : str or SegmentedSFT
        Path to a trk file or SegmentedSFT

    bundle : str
        The name of a bundle to select from the trk metadata.

    colors : dict or list
        If this is a dict, keys are bundle names and values are RGB tuples.
        If this is a list, each item is an RGB tuple. Defaults to a list
        with Tableau 20 RGB values

    n_points : int or None
        n_points to resample streamlines to before plotting. If None, no
        resampling is done.

    n_sls_viz : int
        Number of streamlines to randomly select if plotting
        all bundles. Selections will be proportional to the original number of
        streamlines per bundle.
        Default: 3600
    n_sls_min : int
        Minimun number of streamlines to display per bundle.
        Default: 75

    Returns
    -------
    Statefule Tractogram streamlines, RGB numpy array, str
    """
    viz_logger.info("Loading Stateful Tractogram...")
    if isinstance(trk_file, str):
        seg_sft = aus.SegmentedSFT.fromfile(trk_file)
    else:
        seg_sft = trk_file

    if colors is None:
        colors = gen_color_dict(seg_sft.bundle_names)

    seg_sft.sft.to_vox()
    streamlines = seg_sft.sft.streamlines
    viz_logger.info("Generating colorful lines from tractography...")

    if len(seg_sft.bundle_names) == 1\
            and seg_sft.bundle_names[0] == "whole_brain":
        if isinstance(colors, dict):
            colors = list(colors.values())
        # There are no bundles in here:
        if len(streamlines) > n_sls_viz:
            idx = np.arange(len(streamlines))
            idx = np.random.choice(
                idx, size=n_sls_viz, replace=False)
            streamlines = streamlines[idx]
        if n_points is not None:
            streamlines = dps.set_number_of_points(streamlines, n_points)
        yield streamlines, colors[0], "all_bundles", seg_sft.sft.dimensions
    else:
        if bundle is None:
            # No selection: visualize all of them:
            for bundle_name in seg_sft.bundle_names:
                idx = seg_sft.bundle_idxs[bundle_name]
                if len(idx) == 0:
                    continue
                n_sl_viz = (len(idx) * n_sls_viz) //\
                    len(streamlines)
                n_sl_viz = max(n_sls_min, n_sl_viz)
                if len(idx) > n_sl_viz:
                    idx = np.random.choice(idx, size=n_sl_viz, replace=False)
                these_sls = streamlines[idx]
                if n_points is not None:
                    these_sls = dps.set_number_of_points(these_sls, n_points)
                if isinstance(colors, dict):
                    color = colors[bundle_name]
                else:
                    color = colors[0]
                yield these_sls, color, bundle_name, seg_sft.sft.dimensions
        else:
            these_sls = seg_sft.get_bundle(bundle).streamlines
            if n_points is not None:
                these_sls = dps.set_number_of_points(these_sls, n_points)
            if isinstance(colors, dict):
                color = colors[bundle]
            else:
                color = colors[0]
            yield these_sls, color, bundle, seg_sft.sft.dimensions


def bbox(img):
    img = np.sum(img, axis=-1)
    rows = np.any(img, axis=1)
    cols = np.any(img, axis=0)
    rmin, rmax = np.where(rows)[0][[0, -1]]
    cmin, cmax = np.where(cols)[0][[0, -1]]

    return cmin, rmin, cmax, rmax


def trim(im):
    bg = Image.new(im.mode, im.size, im.getpixel((0, 0)))
    diff = ImageChops.difference(im, bg)
    diff = ImageChops.add(diff, diff, 2.0, -100)
    diff.getbbox
    this_bbox = bbox(diff)
    if this_bbox:
        return im.crop(this_bbox)


def gif_from_pngs(tdir, gif_fname, n_frames,
                  png_fname="tgif", add_zeros=False):
    """
        Helper function
        Stitches together gif from screenshots
    """
    if add_zeros:
        fname_suffix10 = "00000"
        fname_suffix100 = "0000"
        fname_suffix1000 = "000"
    else:
        fname_suffix10 = ""
        fname_suffix100 = ""
        fname_suffix1000 = ""
    angles = []
    n_frame_copies = 60 // n_frames
    for i in range(n_frames):
        if i < 10:
            angle_fname = f"{png_fname}{fname_suffix10}{i}.png"
        elif i < 100:
            angle_fname = f"{png_fname}{fname_suffix100}{i}.png"
        else:
            angle_fname = f"{png_fname}{fname_suffix1000}{i}.png"
        frame = io.imread(op.join(tdir, angle_fname))
        for j in range(n_frame_copies):
            angles.append(frame)

    io.mimsave(gif_fname, angles)


def prepare_roi(roi, affine_or_mapping, static_img,
                roi_affine, static_affine, reg_template):
    """
    Load the ROI
    Possibly perform a transformation on an ROI
    Helper function

    Parameters
    ----------
    roi : str or Nifti1Image
        The ROI information.
        If str, ROI will be loaded using the str as a path.

    affine_or_mapping : ndarray, Nifti1Image, or str
       An affine transformation or mapping to apply to the ROI before
       visualization. Default: no transform.

    static_img: str or Nifti1Image
        Template to resample roi to.

    roi_affine: ndarray

    static_affine: ndarray

    reg_template: str or Nifti1Image
        Template to use for registration.

    Returns
    -------
    ndarray
    """
    viz_logger.info("Preparing ROI...")
    if not isinstance(roi, np.ndarray):
        if isinstance(roi, str):
            roi = nib.load(roi).get_fdata()
        else:
            roi = roi.get_fdata()

    if affine_or_mapping is not None:
        if isinstance(affine_or_mapping, np.ndarray):
            # This is an affine:
            if (static_img is None or roi_affine is None
                    or static_affine is None):
                raise ValueError("If using an affine to transform an ROI, "
                                 "need to also specify all of the following",
                                 "inputs: `static_img`, `roi_affine`, ",
                                 "`static_affine`")
            roi = resample(roi, static_img, roi_affine,
                           static_affine).get_fdata()
        else:
            # Assume it is  a mapping:
            if (isinstance(affine_or_mapping, str)
                    or isinstance(affine_or_mapping, nib.Nifti1Image)):
                if reg_template is None or static_img is None:
                    raise ValueError(
                        "If using a mapping to transform an ROI, need to ",
                        "also specify all of the following inputs: ",
                        "`reg_template`, `static_img`")
                affine_or_mapping = reg.read_mapping(affine_or_mapping,
                                                     static_img,
                                                     reg_template)

            roi = auv.transform_inverse_roi(
                roi,
                affine_or_mapping).astype(bool)
    return roi


def load_volume(volume):
    """
    Load a volume
    Helper function

    Parameters
    ----------
    volume : ndarray or str
        3d volume to load.
        If string, it is used as a file path.
        If it is an ndarray, it is simply returned.

    Returns
    -------
    ndarray
    """
    viz_logger.info("Loading Volume...")
    if isinstance(volume, str):
        return nib.load(volume).get_fdata()
    else:
        return volume


class Viz:
    def __init__(self,
                 backend="fury"):
        """
        Set up visualization preferences.

        Parameters
        ----------
            backend : str, optional
                Should be either "fury" or "plotly".
                Default: "fury"
        """
        self.backend = backend
        if "fury" in backend:
            try:
                import AFQ.viz.fury_backend
            except ImportError:
                raise ImportError(viz_import_msg_error("fury"))
            self.visualize_bundles = AFQ.viz.fury_backend.visualize_bundles
            self.visualize_roi = AFQ.viz.fury_backend.visualize_roi
            self.visualize_volume = AFQ.viz.fury_backend.visualize_volume
            self.create_gif = AFQ.viz.fury_backend.create_gif
        elif "plotly" in backend:
            try:
                import AFQ.viz.plotly_backend
            except ImportError:
                raise ImportError(viz_import_msg_error("plotly"))
            self.visualize_bundles = AFQ.viz.plotly_backend.visualize_bundles
            self.visualize_roi = AFQ.viz.plotly_backend.visualize_roi
            self.visualize_volume = AFQ.viz.plotly_backend.visualize_volume
            self.create_gif = AFQ.viz.plotly_backend.create_gif
            self.single_bundle_viz = AFQ.viz.plotly_backend.single_bundle_viz
        else:
            raise TypeError("Visualization backend contain"
                            + " either 'plotly' or 'fury'. "
                            + "It is currently set to %s"
                            % backend)
