#!/usr/bin/env python

"""

==============================================================================
                  CERN@school Mafalda-to-Pixelman Converter
==============================================================================

Python version of the Mafalda (ROOT) to Pixelman converter.

FIXME: Plenty!

"""

# Import the code needed to manage files.
import sys, os, inspect, glob

# Import the argument parser functionality.
import argparse

# Import the required ROOT libraries - hurrah for PyROOT!
from ROOT import TFile, gSystem
from ROOT import TH1D, TH2D

# ...for the colour scaling.
from matplotlib import colorbar, colors

# ...for the plotting.
import matplotlib.pyplot as plt

# ...for nice text.
from matplotlib import rc

# Uncomment to use LaTeX for the plot text.
rc('font',**{'family':'serif','serif':['Computer Modern']})
rc('text', usetex=True)

import numpy as np

# Load in the (skeleton) Frame class - a bare-minimum class that
# provides the ROOT file format interface.
gSystem.Load('Frame_C')

#
# The main event!
#
if __name__ == "__main__":

    print("==============================================")
    print("     CERN@school Mf2Px-converter (Python)     ")
    print("==============================================")

    dbg = False

    # Get the datafile paths from the command line.
    parser = argparse.ArgumentParser()
    parser.add_argument("inputFilePath", help="The path of the ROOT Timepix data file.")
    parser.add_argument("outputPath",    help="The path for the output files."         )
    parser.add_argument("startFrame",    help="The frame number to start from."        )
    parser.add_argument("numFrames",     help="The number of frames to process."       )
    args = parser.parse_args()

    # Process the input arguments.

    ## The full path of the input Mafalda ROOT file to be analysed.
    datapath = args.inputFilePath

    ## The full path for the output.
    outputpath = args.outputPath

    ## The start frame.
    startframe = int(args.startFrame)

    ## The number of frames to process.
    numframes = int(args.numFrames)

    makeimage = False

    ## The ROOT file itself.
    f = TFile(datapath)

    ## The TTree containing the data.
    chain = f.Get('MPXTree')

    ## The number of frames in the file.
    nframes = chain.GetEntriesFast()

    if startframe > nframes:
        raise IOError("* ERROR: start frame is greater than the number of frames present.")

    # Update the user.
    if dbg:
        print
        print("* Input path is                 : '%s'" % (datapath))
        print("* The output is being written to: '%s/'" % (outputpath))
        print("*")
        print("* Start frame            = %d" % (startframe))
        print("* Number of frames       = %d" % (numframes))
        print("*")
        print("* Number of frames found = %d" % (nframes))

    # Create the mpl Figure and FigCanvas objects.
    # 5x4 inches, 100 dots-per-inch
    #
    dpi = 72

    ## The colour map for the pixel counts.
    cmap = plt.cm.hot # "Hot" used.

    # Loop over the frames in the file.
    #for fn in range(nframes):
    for fn in range(startframe, startframe + numframes + 1):

        # Load the TTree.
        ientry = chain.LoadTree(fn)

        # Copy the entry into memory.
        nb = chain.GetEntry(fn)

        if nb == 0:
            break

        ## The frame metadata (for the DSC file).
        d = ""

        d += "A000000001\n"
        d += "[F0]\n"
        d += "Type=i16 "

        # The payload format.
        fmt = int(chain.FramesData.GetPayloadFormat())
        #print(fmt)
        if fmt == 4114 or fmt == 16384:
            d += "[X,Y,C]"
        else:
            raise IOError("* ERROR: unknown payload format!!!!!!!!!!")
        d += " width=%d height=%d\n" % (chain.FramesData.GetFrameWidth(), chain.FramesData.GetFrameHeight())
        d += "\"Acq mode\" (\"Acquisition mode\"):\n"
        d += "i32[1]\n"
        d += "%d\n" % (chain.FramesData.GetAcqMode())
        d += "\n"
        d += "\"Acq time\" (\"Acquisition time [s]\"):\n"
        d += "double[1]\n"
        d += "%f\n" % (chain.FramesData.GetAcqTime())
        d += "\n"
        d += "\"ChipboardID\" (\"Medipix or chipboard ID\"):\n"
        d += "uchar[10]\n"
        d += "%s\n" % (str(chain.FramesData.GetChipboardID()).strip())
        d += "\n"
        customname = str(chain.FramesData.GetCustomName()).strip()
        if customname != "":
            d += "\"Custom name\" (\"Custom name\"):\n"
            d += "uchar[7]\n"
            d += "%s\n" % (customname)
            d += "\n"
        d += "\"DACs\" (\"DACs values of all chips\"):\n"
        d += "u16[14]\n"

        # The DACs.
        dacs = chain.FramesData.GetDACs()
        for dac in dacs:
            d += "%d " % (dac)
        d += "\n"

        d += "\n"
        d += "\"Firmware\" (\"Firmware version\"):\n"
        d += "char[64]\n"
        d += "%s\n" % (chain.FramesData.GetFirmware())
        d += "\n"
        d += "\"HV\" (\"Bias Voltage [V]\"):\n"
        d += "double[1]\n"
        d += "%f\n" % (chain.FramesData.GetHV())
        d += "\n"
        d += "\"Hw timer\" (\"Hw timer mode\"):\n"
        d += "i32[1]\n"
        d += "%d\n" % (chain.FramesData.GetHwTimerMode())
        d += "\n"
        d += "\"Interface\" (\"Medipix interface\"):\n"
        d += "uchar[4]\n"
        d += "%s\n" % (chain.FramesData.GetInterface())
        d += "\n"
        d += "\"Mpx clock\" (\"Medipix clock [MHz]\"):\n"
        d += "double[1]\n"
        d += "%5.2f\n" % (chain.FramesData.GetMpxClock())
        d += "\n"
        d += "\"Mpx type\" (\"Medipix type (1-2.1, 2-MXR, 3-TPX)\"):\n"
        d += "i32[1]\n"
        d += "%d\n" % (chain.FramesData.GetMpxType())
        d += "\n"
        d += "\"Pixelman version\" (\"Pixelman version\"):\n"
        d += "uchar[6]\n"
        d += "%s\n" % (chain.FramesData.GetPixelmanVersion())
        d += "\n"
        d += "\"Polarity\" (\"Detector polarity (0 negative, 1 positive)\"):\n"
        d += "i32[1]\n"
        d += "%d\n" % (chain.FramesData.GetPolarity())
        d += "\n"
        d += "\"Start time\" (\"Acquisition start time\"):\n"
        d += "double[1]\n"
        d += "%f\n" % (chain.FramesData.GetStartTime())
        d += "\n"
        d += "\"Start time (string)\" (\"Acquisition start time (string)\"):\n"
        d += "char[64]\n"
        d += "%s\n" % (chain.FramesData.GetStartTimeS())
        d += "\n"
        d += "\"Timepix clock\" (\"Timepix clock (0-3: 10MHz, 20MHz, 40MHz, 80MHz)\"):\n"
        d += "byte[1]\n"
        d += "0\n"

        ## The hit pixels in the frame.
        pxs = chain.FramesData.GetFrameXC()

        print("* Frame %3d - number of raw pixels found      = %6d" % (fn, len(pxs)))

        ## Dictionary for the pixels.
        pixels = {}

        ## The masked pixels from the ROOT file.
        maskedPixels = chain.FramesData.GetLVL1()

        ## The pixel mask dictionary.
        mask = {}
        #
        for i in maskedPixels:
             if i.second == 1:
                 mask[i.first] = i.second

        ## The data file string to write to the file.
        datw = ""

        # Loop over the pixels and extract the coordinates.
        for px in pxs:

            if px.first not in mask.keys():
                X = px.first
                x = X%256
                y = X/256
                C = px.second

                # Add the pixel to the pixel dictionary.
                pixels[X] = C

                # Add the pixel to the data file.
                datw += "%d\t%d\t%d\n" % (x,y,C)

        print("*              - number of unmasked pixels found = %6d" % (len(pixels)))
        print("*")

        # Create the data file.
        datf = open(outputpath+"/data%05d.txt" % (fn),"w")

        # Write the data file.
        datf.write(datw)

        # Create the DSC file.
        dscf = open(outputpath+"/data%05d.txt.dsc" % (fn), "w")

        # Write to the DSC file.
        dscf.write(d)

        # Close the output text files.
        datf.close()
        dscf.close()

        # Create the frame image.
        if makeimage:
            ## The current maximum count value.
            cmax = 500 #max(pixels.values())

            # The figure and array of axes for the frame display.
            mapfig, mapaxarr = plt.subplots(1, 1, figsize=(8.0,8.0), dpi=dpi, frameon=False)
            mapfig.subplots_adjust(0.01,0.01,0.99,0.99)

            # Remove the axis labels.
            mapaxarr.xaxis.set_visible(False)
            mapaxarr.yaxis.set_visible(False)
            mapaxarr.set_xlim([0,256])
            mapaxarr.set_ylim([0,256])

            # Clear the display.
            mapaxarr.clear()

            # Set the background colour.
            mapaxarr.set_axis_bgcolor('#82bcff')

            # Add a grid to the plot.
            mapaxarr.grid(1)

            # Set the plot title.
            #mapaxarr.set_title("Frame %05d" % (fn))

            # Loop over the pixels in the pixel dictionary.
            for X, C in pixels.iteritems():
                ## The pixel x value.
                x = X%256;

                ## The pixel y value.
                y = X/256
                #print("* DEBUG: Pixel found at (%3d, %3d) with C = %d" % (x,y,C))
                #if C >= threshold:
                color = float(C) / float(cmax)
                #color = ((float(C) * np.e) / float(cmax) + 4.5) / 5.5
                #print("* DEBUG: C = %d - setting color to %f" % (C, color))

                # Plot the pixel.
                mapaxarr.add_patch(plt.Rectangle((x,y),1,1, edgecolor=cmap(color), facecolor=cmap(color)))

            mapfig.savefig(outputpath + "/frame%05d.png" % (fn))

    # Close the data file.
    f.Close()