# standard libraries imports
from __future__ import print_function
import os
import itertools

# external libraries imports
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

# fullrmc library imports
from fullrmc.Engine import Engine

# create figure
FIG = plt.figure()
FIG.patch.set_facecolor('white')
grid = gridspec.GridSpec(nrows=2, ncols=2)
grid.update(left=0.1, bottom=0.1, right=0.95, wspace=0.05)

totalAx = plt.subplot(grid[0, :])
lowRAx  = plt.subplot(grid[1, 0])
highRAx = plt.subplot(grid[1, 1])


# set axis ticks
highRAx.get_yaxis().set_ticks([])
lowRAx.get_yaxis().set_ticks([])

# Create plotting styles
colors = ["b",'g','r','c','m','y']
markers = ["",'.','+','^','|']
INTRA_STYLES = [r[0] + r[1]for r in itertools.product(['--'], colors)]
INTRA_STYLES = [r[0] + r[1]for r in itertools.product(markers, INTRA_STYLES)]
INTER_STYLES = [r[0] + r[1]for r in itertools.product(['-'], colors)]
INTER_STYLES = [r[0] + r[1]for r in itertools.product(markers, INTER_STYLES)]

# dirname
DIR_PATH = os.path.dirname( os.path.realpath(__file__) )
engineFilePath = os.path.join(DIR_PATH, "thf_engine.rmc")

# load
ENGINE = Engine(path=None)
result, mes = ENGINE.is_engine(engineFilePath, mes=True)
if result:
    ENGINE = ENGINE.load(engineFilePath)
    PDF = ENGINE.constraints[0]
else:
    print(mes)


# plot total
output = PDF.get_constraint_value()
totalAx.plot(PDF.experimentalDistances, ENGINE.constraints[0].experimentalPDF, 'ro', label="observed", markersize=7.5, markevery=1 )
lowRAx.plot(PDF.experimentalDistances, ENGINE.constraints[0].experimentalPDF, 'ro', label="observed", markersize=7.5, markevery=1 )
highRAx.plot(PDF.experimentalDistances, ENGINE.constraints[0].experimentalPDF, 'ro', label="observed", markersize=7.5, markevery=1 )
totalAx.plot(PDF.shellCenters, output["pcf"], 'k', linewidth=3.0,  markevery=25, label="total" )
lowRAx.plot(PDF.shellCenters, output["pcf"], 'k', linewidth=3.0,  markevery=25, label="total" )
highRAx.plot(PDF.shellCenters, output["pcf"], 'k', linewidth=3.0,  markevery=25, label="total" )

# plot inter and intra
#intraStyleIndex = 0
#interStyleIndex = 0
#for key in output:
#    val = output[key]
#    if key in ("pdf_total", "pdf"):
#        continue
#    elif "intra" in key:
#        totalAx.plot(PDF.shellCenters, val, INTRA_STYLES[intraStyleIndex], markevery=5, label=key.split("rdf_")[1] )
#        lowRAx.plot(PDF.shellCenters, val, INTRA_STYLES[intraStyleIndex], markevery=5, label=key.split("rdf_")[1] )
#        highRAx.plot(PDF.shellCenters, val, INTRA_STYLES[intraStyleIndex], markevery=5, label=key.split("rdf_")[1] )
#        intraStyleIndex+=1
#    elif "inter" in key:
#        totalAx.plot(PDF.shellCenters, val, INTER_STYLES[interStyleIndex], markevery=5, label=key.split("rdf_")[1] )
#        lowRAx.plot(PDF.shellCenters, val, INTRA_STYLES[intraStyleIndex], markevery=5, label=key.split("rdf_")[1] )
#        highRAx.plot(PDF.shellCenters, val, INTRA_STYLES[intraStyleIndex], markevery=5, label=key.split("rdf_")[1] )
#        interStyleIndex+=1

# set legend
totalAx.legend(ncol=2, frameon=False, fontsize=20)#, loc=(1.05,-1.25))
totalAx.set_xlim([0,16])
totalAx.set_ylim([-1,6.2])

# remove y ticks labels
lowRAx.set_yticklabels([])
highRAx.set_yticklabels([])

# set x limits
lowRAx.relim()
highRAx.relim()
lowRAx.set_xlim([0.5,3])
lowRAx.set_ylim([-1, 6.2])
highRAx.set_xlim([3,16])
highRAx.set_ylim([0.25, 1.25])

# add title
plt.suptitle("$stdErr=%.6f$"%ENGINE.totalStandardError, fontsize=20)

# show plot
plt.show()
