#!/usr/bin/env python3

# Imports
def cm2inch(value):
    return value/2.54

import matplotlib.pyplot as plt
import matplotlib.cm as cm
import pandas as pd
import numpy as np
import seaborn.apionly as sns


niceblue = '#64B5CD'
nicered = '#C44E52'
niceyellow = '#CCB974'
plt.style.use('custom')

A = pd.read_csv('janus_quench.csv')
A = A[A['observable'] == "Energy"]
A = A[A['T'] == 0.9]
A = A[A['L'] == 80]
print("Replicas para el quench:", len(np.unique(A['Jhash'])))
A = A.groupby(['mc']).mean().reset_index()
A = A[A['mc'] <= 1e8]

B = pd.read_csv('janus_continuation.csv')
B = B[B['observable'] == "Energy"]
B = B[B['T1'] == 0.9]
B = B[B['T2'] == 0.5]
B = B[B['tw1'] == 1e8]
print("Replicas para el cont:", len(np.unique(B['Jhash'])))
B = B.groupby(['mc']).mean().reset_index()


C = pd.read_csv('janus_finish.csv')
C = C[C['observable'] == "Energy"]
C = C[C['T1'] == 0.9]
C = C[C['T2'] == 0.5]
C = C[C['tw1'] == 1e8]
print("Replicas para el finish:", len(np.unique(C['Jhash'])))
C = C.groupby(['mc']).mean().reset_index()

# Just the 1D plots
fig = plt.figure(figsize=(cm2inch(15),cm2inch(5)))
ax = fig.add_subplot(111)

alldata = pd.concat([A,B,C], ignore_index=True)

ax.scatter(alldata['mc']*1e-8, alldata['value'], s=2, color=niceblue)

# ax.set_xscale('log')
ax.set_xlabel('$MC [×10^8]$')
ax.set_ylabel('[E]')

ax.text(0.5, -1.7, 'T=0.9', fontsize=11, color='gray', horizontalalignment='center')
ax.text(1.5, -1.7, 'T=0.5', fontsize=11, color='gray', horizontalalignment='center')
ax.text(2.5, -1.7, 'T=0.9', fontsize=11, color='gray', horizontalalignment='center')


fig.savefig('energyprotocol.pdf')

# For the beamer
ax.grid(b=False)
fig.set_size_inches(cm2inch(10),cm2inch(5))
fig.savefig('energyprotocol_beamer.pdf')
