import glob
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.ticker as mtick
from matplotlib.colors import Normalize

import sys
from pathlib import Path
import numpy as np
import pandas as pd


# -------------------------------------------------------------------------------
# This file plots data generated by MASIC.exe or MASIC_Console.exe, saving plots as PNG files
#
# Written by Matthew Monroe for the Department of Energy (PNNL, Richland, WA)
# Program started in 2020
#
# E-mail: matthew.monroe@pnnl.gov or proteomics@pnnl.gov
# Website: https://github.com/PNNL-Comp-Mass-Spec/ or https://panomics.pnnl.gov/ or https://www.pnnl.gov/integrative-omics
# -------------------------------------------------------------------------------

# When this script is invoked from C#, if text sent to print() has unicode characters, an error occurs
# For example, for the micron symbol, we see:
#   UnicodeEncodeError: 'charmap' codec can't encode character '\u03bc' in position 45: character maps to <undefined>

# To fix this, we can either force the stdout and stderr encoding to utf-8, like this:
#   import codecs
#   if sys.stdout.encoding != 'UTF-8':
#     sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, 'strict')
#   if sys.stderr.encoding != 'UTF-8':
#     sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, 'strict')

# Or we can use the uprint function shown below

# Update the default font
plt.rcParams['font.family'] = ['Arial', 'FreeSans', 'DejaVu Sans']


def process_file(dataFilePath):
    dataFile = Path(dataFilePath)
    if not dataFile.is_file():
        uprint('\nError, file not found: ' + dataFilePath)
        return

    uprint('Input: ', dataFile)
    data, plotOptions, columnOptions = read_file(dataFile)

    outputFilePath = str(dataFile.with_suffix('.png'))

    if 'PlotType' not in plotOptions:
        print('Invalid plot options; PlotType must be defined')
        print(plotOptions)
        return

    if plotOptions['PlotType'] == 'XY':
        # Histogram (2D plot)
        uprint('Output:', outputFilePath)
        print()
        uprint('Plot "' + data.columns[0] + '" vs. "' + data.columns[1] + '"')
        print("  {:,}".format(len(data.index)) + ' data points')
        print()
        plot_histogram(outputFilePath, data.columns, data[data.columns[0]], data[data.columns[1]],
                       plotOptions['Title'], plotOptions['BottomRight'])
        return

    if plotOptions['PlotType'] == 'BarChart':
        # Bar chart
        isPercentages = string_to_bool(plotOptions['Percentages'])
        
        uprint('Output:', outputFilePath)
        print()
        uprint('Plot "' + data.columns[0] + '" vs. "' + data.columns[1])
        print("  {:,}".format(len(data.index)) + ' data points')
        print()
        plot_bar_chart(outputFilePath, data.columns, data[data.columns[0]], data[data.columns[1]],
                       plotOptions['Title'], isPercentages, plotOptions['BottomRight'], plotOptions['BottomLeft'])
        return

    if plotOptions['PlotType'] == 'BoxPlot':
        # Box and whisker plot
        uprint('Output:', outputFilePath)
        print()
        uprint('Plot "' + data.columns[0] + '" vs. "' + data.columns[1])
        print("  {:,}".format(len(data.index)) + ' data points')
        print()
        plot_box_chart(outputFilePath, data.columns, data[data.columns[0]], data[data.columns[1]],
                       plotOptions['Title'], plotOptions['BottomRight'], plotOptions['BottomLeft'])
        return

    print('Unsupported plot type: ' + plotOptions['PlotType'])
    print(data.columns)


def string_to_bool(value):
    if value.lower() == 'true':
         return True
    elif value.lower() == 'false':
         return False
    else:
         raise ValueError("Cannot convert {} to a bool".format(value))


def parse_metadata(plotOption):
    return {m.split('=')[0]: m.split('=')[1] for m in plotOption.split(';')}


def read_file(fpath):
    data = pd.read_csv(fpath, sep='\t', skiprows=2, header=0)
    with open(fpath, 'r') as f:
        # The first line has the plot type, plot title, etc., formatted as:
        # [PlotType=XY;Title=PlotTitle;Percentages=True;BottomLeft=Annotation1;BottomRight=Annotation2;]
        plotOptionData = f.readline().split('[')[1].split(']')[0]
        plotOptions = parse_metadata(plotOptionData)

        # The second line has column options
        # These are semicolon separated key/value pairs for each column, with options for each column separated by a tab
        # At present, the code does not use these column options
        columnOptionData = f.readline().split('\t')
        columnOptions = [parse_metadata(colOption) for colOption in columnOptionData]

        return data, plotOptions, columnOptions


def set_title_and_labels(ax, baseFontSize, plotType, title, isPercentages,
                         xDataMin, xDataMax, 
                         xAxisLabel, yAxisLabel, yAxisFormatString,
                         r_label, l_label):

    MAX_TITLE_LENGTH = 68

    if len(title) > MAX_TITLE_LENGTH:
        # Trim out the middle portion of the title
        # Find the first space
        spaceIndex = title.find(' ')
        if spaceIndex > 0:
            titleSuffix = title[spaceIndex:]
            charsToKeep = MAX_TITLE_LENGTH - len(titleSuffix)
            if charsToKeep > 0:
                titleToUse = title[0:charsToKeep] + titleSuffix
            else:
                titleToUse = title[-MAX_TITLE_LENGTH:]

        else:
            titleToUse = title[-MAX_TITLE_LENGTH:]
    else:
        titleToUse = title

    plt.title(titleToUse, fontsize=baseFontSize + 1)

    # Assure that the X axis minimum is not negative
    xmin, xmax = plt.xlim()

    if plotType == "Histogram" and xmin < 0:
        plt.xlim(xmin=0)

    # When plotting a histogram or a bar chart of percentages, fix the Y axis minimum at 0 and add 2% padding above the Y axis maximum
    # Otherwise, assure that the Y range isn't too small
    ymin, ymax = plt.ylim()

    if plotType == "Histogram":
        ymin = 0
        ymax = ymax * 1.02
    elif plotType == "BarChart":
        if isPercentages:
            # Bar Chart of values between 0 and 100
            # Scale the y axis from 0 to 102 (to give a little whitespace above bars that are at 100%)
            ymin = 0
            ymax = 102
        else:
            ymin = 0
            ymax = ymax

    plt.ylim(ymin=ymin, ymax=ymax)

    if plotType == "Histogram":
        # Set the X axis maximum to the max X value (in other words, we don't want any padding)
        # However, if there is only one X data point, do add some padding
        if xDataMin == xDataMax:
            plt.xlim(xmin=xDataMax - 1)
            plt.xlim(xmax=xDataMax + 1)
        else:
            plt.xlim(xmax=xDataMax)

    if len(xAxisLabel) > 0:
        plt.xlabel(xAxisLabel, fontsize=baseFontSize)

    if len(yAxisLabel) > 0:
        plt.ylabel(yAxisLabel, fontsize=baseFontSize)

    if plotType == "Histogram":
        plt.xticks(fontsize=baseFontSize - 2)

    plt.yticks(fontsize=baseFontSize - 2)
    ax.yaxis.set_major_formatter(mtick.FormatStrFormatter(yAxisFormatString))

    # Define the minor tick locator for the y axis, only if the y-axis is not log scaled
    if type(ax.yaxis._scale) != mpl.scale.LogScale:
        ax.yaxis.set_minor_locator(mtick.AutoMinorLocator())

    # If the x axis range is 5 or less, assure that the minimum distance between major tick marks is at least 1
    if not (xDataMin == 0 and xDataMax == 0):
        if xDataMax - xDataMin <= 5:
            ax.xaxis.set_major_locator(mtick.MultipleLocator(1))

    if plotType == "Histogram":
        ax.xaxis.set_major_formatter(mtick.FuncFormatter(lambda x, p: format(int(x), ',')))
        ax.xaxis.set_minor_locator(mtick.AutoMinorLocator())

    plt.gcf().text(0.88, 0.02, r_label, fontsize=baseFontSize - 1)

    if len(l_label) > 0:
        plt.gcf().text(0.01, 0.02, l_label, fontsize=baseFontSize - 1)


def plot_histogram(outputFilePath, columnNames, xValues, yValues, title, r_label):
    fig = plt.figure(figsize=(8.5333, 5), dpi=120)
    ax = fig.add_subplot(111, facecolor='whitesmoke')

    baseFontSize = 12

    if len(yValues) == 1:
        print('Single point plot')
        ax.plot(xValues, yValues, linewidth=1, markersize=4, marker='o', color='darkblue')
    else:
        ax.plot(xValues, yValues, linewidth=1, color='black')

    # X axis
    xAxisLabel = columnNames[0]

    # Y axis
    yAxisLabel = columnNames[1]

    ymin, ymax = plt.ylim()

    if ymax < 10000:
        yAxisFormatString = '%d'
    else:
        yAxisFormatString = '%.2e'
        
    set_title_and_labels(ax, baseFontSize, "Histogram", title, False,
                         np.min(xValues), np.max(xValues),
                         xAxisLabel, yAxisLabel, yAxisFormatString, 
                         r_label, '')

    plt.tight_layout()

    plt.savefig(outputFilePath)

    print('XY plot created')

    # Uncomment to view the plot with an interactive GUI
    # plt.show()


def generate_bar_chart(columnNames, xLabels, barHeights, title, isPercentages, r_label, l_label):
    fig = plt.figure(figsize=(8.5333, 5.8333), dpi=120)
    ax = fig.add_subplot(111, facecolor='whitesmoke')

    baseFontSize = 12

    # dataMin = np.min(barHeights)
    # dataMax = np.max(barHeights)

    # Define the range of values that we could see
    # This is used when normalizing the bar heights to determine the bar color
    dataMin = 0
    dataMax = 100

    cmap = plt.cm.get_cmap('inferno_r')

    dataNormalizer = Normalize(vmin=dataMin, vmax=dataMax)
    normList = dataNormalizer(barHeights.tolist())

    # Uncomment to debug
    # print('Min:     ' + "{:10.2f}".format(dataMin))
    # print('Max:     ' + "{:10.2f}".format(dataMax))
    # print('Range:   ' + "{:10.2f}".format(dataRange))

    # Construct a list of bar positions
    xPos = [i for i, _ in enumerate(xLabels)]

    ax.bar(xPos, barHeights, color=cmap(normList), edgecolor='black')

    # Leave the x axis label blank, since this is a bar chart
    xAxisLabel = ""

    # Y axis is counts
    yAxisLabel = columnNames[1]

    yAxisFormatString = '%d'
        
    set_title_and_labels(ax, baseFontSize, "BarChart", title, isPercentages,
                         0, 0,
                         xAxisLabel, yAxisLabel, yAxisFormatString, 
                         r_label, l_label)

    # Add bar labels (rotated)
    plt.xticks(xPos, xLabels)
    plt.setp(plt.gca().get_xticklabels(), rotation=-30, horizontalalignment='left')

    plt.tight_layout()

    return plt


def generate_box_plot(columnNames, xLabels, xData, title, r_label, l_label):
    """
    Generate a box and whiskers plot
    :param columnNames: Column names, tracking the X and Y axis titles
    :param xLabels: X axis labels
    :param xData: Y axis data, with the same number of rows as xLabels; each row is a comma-separated list of values
    :param title: Plot title
    :param r_label: Optional text for the bottom right of the plot
    :param l_label: Optional text for the bottom left of the plot
    :return:
    """

    fig = plt.figure(figsize=(8.5333, 5.8333), dpi=120)
    ax = fig.add_subplot(111, facecolor='whitesmoke')

    baseFontSize = 12

    # Construct a list of bar positions
    xPos = [i+1 for i, _ in enumerate(xLabels)]

    # Transform xData into a sequence of vectors
    boxData = []

    for intensityList in xData:
        if isinstance(intensityList, str):
            # List of floats; split on commas
            # This also works if the column has a single float or single integer
            intensities = [float(i) for i in intensityList.split(',')]
        else:
            # This line had an empty second column in the input ifle
            intensities = []

        boxData.append(intensities)

    boxprops = dict(linestyle='-', linewidth=1.5, color='black')
    flierprops = dict(marker='o', markersize=3,
                      linestyle='none')
    medianprops = dict(linestyle='-', linewidth=2, color='black')

    bp = ax.boxplot(boxData, boxprops=boxprops, flierprops=flierprops, medianprops=medianprops, patch_artist=True)

    #for element in ['boxes', 'whiskers', 'fliers', 'means', 'medians', 'caps']:
    #    plt.setp(bp[element], color='black')

    # This could be used to give each box a different color
    for patch in bp['boxes']:
        patch.set(facecolor='lightskyblue')

    # Leave the x axis label blank, since this is a bar chart
    xAxisLabel = ""

    # Y axis is counts
    yAxisLabel = columnNames[1]

    yAxisFormatString = '%.0e'

    # Use logarithmic scaling for the y-axis
    ax.set_yscale('log')

    set_title_and_labels(ax, baseFontSize, "BoxPlot", title, False,
                         0, 0,
                         xAxisLabel, yAxisLabel, yAxisFormatString, 
                         r_label, l_label)

    # Add bar labels (rotated)
    plt.xticks(xPos, xLabels)
    plt.setp(plt.gca().get_xticklabels(), rotation=-30, horizontalalignment='left')

    plt.tight_layout()

    return plt


def plot_bar_chart(outputFilePath, columnNames, xLabels, barHeights, title, isPercentages, r_label, l_label):
    barChart = generate_bar_chart(columnNames, xLabels, barHeights, title, isPercentages, r_label, l_label)
    barChart.savefig(outputFilePath)
    print('Bar chart created')


def plot_box_chart(outputFilePath, columnNames, xLabels, xData, title, r_label, l_label):
    boxPlot = generate_box_plot(columnNames, xLabels, xData, title, r_label, l_label)
    boxPlot.savefig(outputFilePath)
    print('Box plot created')


def uprint(*objects, sep=' ', end='\n', file=sys.stdout):
    enc = file.encoding
    
    if enc.lower() == 'utf-8':
        print(*objects, sep=sep, end=end, file=file)
    else:
        f = lambda obj: str(obj).encode(enc, errors='backslashreplace').decode(enc)
        print(*map(f, objects), sep=sep, end=end, file=file)

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print('\nError: please enter the file name to process (wildcards are supported)')
        exit()

    fileNameMatchSpec = sys.argv[1]

    filesProcessed = 0
    for matchedFile in glob.glob(fileNameMatchSpec):
        process_file(matchedFile)
        filesProcessed += 1

    if filesProcessed == 0:
        print('\nError: no files match:\n' + fileNameMatchSpec)

# Can list fonts with the following
# import matplotlib.font_manager
# from matplotlib.font_manager import findfont, FontProperties
# fonts = matplotlib.font_manager.findSystemFonts(fontpaths=None, fontext='ttf')
# for font in sorted(fonts):
# print(font)

# View details on a specific font with
# font = findfont(FontProperties(family=['arial']))
# print (font)
# exit()
