#!/usr/bin/env python3

from re import match, split
import numpy as np
from scipy.interpolate import interp1d
from pathlib import Path
from ROOT import gStyle, TGraph, TMultiGraph, TCanvas, TLegend
from ROOT import SetOwnership, gPad, gROOT

gROOT.Reset()

datadir = str(Path(__file__).resolve().parent) + "/"

mg = TMultiGraph()
gr0 = TGraph()
gr1 = TGraph()
gr2 = TGraph()
mg.Add(gr0, "l")
mg.Add(gr1, "l")
mg.Add(gr2, "l")


def depict_graph(gr, color, sp):
    gr.SetLineColor(color)
    gr.SetLineWidth(3)
    gr.SetLineStyle(1)
    for i in range(len(sp)):
        gr.SetPoint(i, energy[i], sp[i])

mass_number = 177

##############################
# LET by PHITS
##############################
energy = []
sp_ele = []
sp_nucl = []
sp_tot = []

f0 = open(datadir + "srim_result.txt")
for row in f0:
    if match(r"\s\s*\d", row):
        row = row.strip()
        item = row.split()
        # print(item)
        if item[1] == "keV":
            energy.append(float(item[0]) / 1e3 / mass_number)
        if item[1] == "MeV":
            energy.append(float(item[0]) / mass_number)
        sp_ele.append(float(item[2]))
        sp_nucl.append(float(item[3]))
        sp_tot.append(float(item[2]) + float(item[3]))

depict_graph(gr0, 1, sp_tot)
depict_graph(gr1, 2, sp_ele)
depict_graph(gr2, 4, sp_nucl)

c0 = TCanvas("c0", "c0", 30, 30, 600, 600)

c0.ToggleEditor()
c0.ToggleEventStatus()
c0.ToggleToolBar()
c0.SetTickx(1)
c0.SetTicky(1)
c0.SetTopMargin(0.05)
c0.SetRightMargin(0.05)
c0.SetLeftMargin(0.17)
c0.SetBottomMargin(0.14)

c0.cd()
gPad.SetLogx(1)
gPad.SetLogy(1)
gStyle.SetLabelFont(42, "xy")
gStyle.SetLabelSize(0.06, "xy")
gStyle.SetLabelOffset(-1e-2, "x")
gStyle.SetLabelOffset(1e-2, "y")
gStyle.SetTitleFont(42, "xy")
gStyle.SetTitleSize(0.06, "xy")
gStyle.SetTitleOffset(1.1, "x")
gStyle.SetTitleOffset(1.3, "y")

mg.Draw("APL")
mg.GetXaxis().SetNdivisions(505)
mg.GetXaxis().SetLimits(9e-5, 2e-1)
mg.GetXaxis().SetTitle("Nuclei Energy [MeV/u]")
mg.GetXaxis().CenterTitle(1)
mg.GetYaxis().SetTitle("LET [MeV cm^{2}/mg]")
mg.GetYaxis().SetNdivisions(505)
mg.GetYaxis().CenterTitle(1)
mg.GetYaxis().SetRangeUser(9e-1, 3e1)

leg0 = TLegend(0.20, 0.60, 0.56, 0.80)
leg0.SetFillStyle(0)
leg0.AddEntry(gr0, "Total", "l")
leg0.AddEntry(gr1, "Electronic", "l")
leg0.AddEntry(gr2, "Nuclear", "l")
leg0.SetTextSize(0.06)
leg0.SetBorderSize(0)
leg0.Draw("SAME")

c0.Modified()
c0.Update()

c0.Print(input("filename = ")+".eps")

c0.Close()

exit()