import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

#Load file
dt=pd.read_csv("sevn_output/output_0.csv")
#Give a look to the columns
print(dt.columns)
#Consider only the final states
dt=dt.drop_duplicates(["ID","name"], keep='last')

#Load evolved file
dte=pd.read_csv("sevn_output/evolved_0.dat",sep='\s+')
#Give a look to the columns
print(dte.columns)
dte=dte.rename(columns={'#ID': 'ID','Mass_0':"Mzams_0", 'Mass_1':"Mzams_1"})
#After change
print(dte.columns)

#Join the two dataset
dt = dt.merge(dte, on=["ID","name"],  how="inner", suffixes=("","_ini") )
# - on: column(s, can be a list of columns) to match during the merge of the two tables. The colum(s) has(have) to be present in both the tables
# - how: type of join to use, see documentation here and the next slide
# - suffixes: columns with the same name in the two tables (not used in on) will be renamed adding these suffixes.
#Give a look to the columns
print(dt.columns)


#Create filter indexes
idx0 = (dt.RemnantType_0==6)
idx1 = (dt.RemnantType_1==6)
idxb0 = idx0  & dt.Semimajor.notnull()
idxb1 = idx1  & dt.Semimajor.notnull()
idxm0 = idxb0 & (dt.GWtime + dt.BWorldtime  <= 14000)
idxm1 = idxb1 & (dt.GWtime + dt.BWorldtime  <= 14000)


#Filter and join masses
AllBH = pd.concat([dt[idx0].Mass_0,dt[idx1].Mass_1])
BoundBH = pd.concat([dt[idxb0].Mass_0,dt[idxb1].Mass_1])
MergingBH = pd.concat([dt[idxm0].Mass_0,dt[idxm1].Mass_1])

#Filter and join initial masses
AllBHzams = pd.concat([dt[idx0].Mzams_0,dt[idx1].Mzams_1])
BoundBHzams = pd.concat([dt[idxb0].Mzams_0,dt[idxb1].Mzams_1])
MergingBHzams = pd.concat([dt[idxm0].Mzams_0,dt[idxm1].Mzams_1])

#Filter and join initial semimajor axis
AllBHa = pd.concat([dt[idx0].a,dt[idx1].a])
BoundBHa = pd.concat([dt[idxb0].a,dt[idxb1].a])
MergingBHa = pd.concat([dt[idxm0].a,dt[idxm1].a])


#Plot
plt.figure(figsize=(10,5))

plt.subplot(1,2,1)
plt.scatter(AllBHzams,AllBH,zorder=1,edgecolor="k",s=30,label="All")
plt.scatter(BoundBHzams,BoundBH,zorder=2,edgecolor="k",s=30, label="Bound")
plt.scatter(MergingBHzams,MergingBH,zorder=3,edgecolor="k",s=30, label="Merging")
plt.plot(np.linspace(0,140),np.linspace(0,140),ls="dashed",c="gray")
plt.xscale("log")
plt.yscale("log")
plt.ylabel("BH mass [M$_\odot$]",fontsize=18)
plt.xlabel("$M\mathrm{zams}$  [M$_\odot$]",fontsize=18)
plt.gca().tick_params(axis='both', which='major', labelsize=18)
plt.legend(fontsize=16)

plt.subplot(1,2,2)
plt.scatter(AllBHa,AllBH,zorder=1,edgecolor="k",s=30,label="All")
plt.scatter(BoundBHa,BoundBH,zorder=2,edgecolor="k",s=30,label="Bound")
plt.scatter(MergingBHa,MergingBH,zorder=3,edgecolor="k",s=30,label="Merging")
plt.xscale("log")
plt.yscale("log")
plt.xlabel("Semimajor initial  [R$_\odot$]",fontsize=18)
plt.ylabel("BH mass [M$_\odot$]",fontsize=18)
plt.gca().tick_params(axis='both', which='major', labelsize=18)



plt.tight_layout()
plt.savefig("analysis3.png")
plt.show()

[(6, 'pandas.read_csv', 'pd.read_csv', (['"""sevn_output/output_0.csv"""'], {}), True, 'import pandas as pd\n'), (13, 'pandas.read_csv', 'pd.read_csv', (['"""sevn_output/evolved_0.dat"""'], {'sep': '"""\\\\s+"""'}), True, 'import pandas as pd\n'), (39, 'pandas.concat', 'pd.concat', (['[dt[idx0].Mass_0, dt[idx1].Mass_1]'], {}), True, 'import pandas as pd\n'), (40, 'pandas.concat', 'pd.concat', (['[dt[idxb0].Mass_0, dt[idxb1].Mass_1]'], {}), True, 'import pandas as pd\n'), (41, 'pandas.concat', 'pd.concat', (['[dt[idxm0].Mass_0, dt[idxm1].Mass_1]'], {}), True, 'import pandas as pd\n'), (44, 'pandas.concat', 'pd.concat', (['[dt[idx0].Mzams_0, dt[idx1].Mzams_1]'], {}), True, 'import pandas as pd\n'), (45, 'pandas.concat', 'pd.concat', (['[dt[idxb0].Mzams_0, dt[idxb1].Mzams_1]'], {}), True, 'import pandas as pd\n'), (46, 'pandas.concat', 'pd.concat', (['[dt[idxm0].Mzams_0, dt[idxm1].Mzams_1]'], {}), True, 'import pandas as pd\n'), (49, 'pandas.concat', 'pd.concat', (['[dt[idx0].a, dt[idx1].a]'], {}), True, 'import pandas as pd\n'), (50, 'pandas.concat', 'pd.concat', (['[dt[idxb0].a, dt[idxb1].a]'], {}), True, 'import pandas as pd\n'), (51, 'pandas.concat', 'pd.concat', (['[dt[idxm0].a, dt[idxm1].a]'], {}), True, 'import pandas as pd\n'), (55, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(10, 5)'}), True, 'import matplotlib.pyplot as plt\n'), (57, 'matplotlib.pyplot.subplot', 'plt.subplot', (['(1)', '(2)', '(1)'], {}), True, 'import matplotlib.pyplot as plt\n'), (58, 'matplotlib.pyplot.scatter', 'plt.scatter', (['AllBHzams', 'AllBH'], {'zorder': '(1)', 'edgecolor': '"""k"""', 's': '(30)', 'label': '"""All"""'}), True, 'import matplotlib.pyplot as plt\n'), (59, 'matplotlib.pyplot.scatter', 'plt.scatter', (['BoundBHzams', 'BoundBH'], {'zorder': '(2)', 'edgecolor': '"""k"""', 's': '(30)', 'label': '"""Bound"""'}), True, 'import matplotlib.pyplot as plt\n'), (60, 'matplotlib.pyplot.scatter', 'plt.scatter', (['MergingBHzams', 'MergingBH'], {'zorder': '(3)', 'edgecolor': '"""k"""', 's': '(30)', 'label': '"""Merging"""'}), True, 'import matplotlib.pyplot as plt\n'), (62, 'matplotlib.pyplot.xscale', 'plt.xscale', (['"""log"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (63, 'matplotlib.pyplot.yscale', 'plt.yscale', (['"""log"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (64, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""BH mass [M$_\\\\odot$]"""'], {'fontsize': '(18)'}), True, 'import matplotlib.pyplot as plt\n'), (65, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""$M\\\\mathrm{zams}$  [M$_\\\\odot$]"""'], {'fontsize': '(18)'}), True, 'import matplotlib.pyplot as plt\n'), (67, 'matplotlib.pyplot.legend', 'plt.legend', ([], {'fontsize': '(16)'}), True, 'import matplotlib.pyplot as plt\n'), (69, 'matplotlib.pyplot.subplot', 'plt.subplot', (['(1)', '(2)', '(2)'], {}), True, 'import matplotlib.pyplot as plt\n'), (70, 'matplotlib.pyplot.scatter', 'plt.scatter', (['AllBHa', 'AllBH'], {'zorder': '(1)', 'edgecolor': '"""k"""', 's': '(30)', 'label': '"""All"""'}), True, 'import matplotlib.pyplot as plt\n'), (71, 'matplotlib.pyplot.scatter', 'plt.scatter', (['BoundBHa', 'BoundBH'], {'zorder': '(2)', 'edgecolor': '"""k"""', 's': '(30)', 'label': '"""Bound"""'}), True, 'import matplotlib.pyplot as plt\n'), (72, 'matplotlib.pyplot.scatter', 'plt.scatter', (['MergingBHa', 'MergingBH'], {'zorder': '(3)', 'edgecolor': '"""k"""', 's': '(30)', 'label': '"""Merging"""'}), True, 'import matplotlib.pyplot as plt\n'), (73, 'matplotlib.pyplot.xscale', 'plt.xscale', (['"""log"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (74, 'matplotlib.pyplot.yscale', 'plt.yscale', (['"""log"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (75, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""Semimajor initial  [R$_\\\\odot$]"""'], {'fontsize': '(18)'}), True, 'import matplotlib.pyplot as plt\n'), (76, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""BH mass [M$_\\\\odot$]"""'], {'fontsize': '(18)'}), True, 'import matplotlib.pyplot as plt\n'), (81, 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (82, 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""analysis3.png"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (83, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (61, 'numpy.linspace', 'np.linspace', (['(0)', '(140)'], {}), True, 'import numpy as np\n'), (61, 'numpy.linspace', 'np.linspace', (['(0)', '(140)'], {}), True, 'import numpy as np\n'), (66, 'matplotlib.pyplot.gca', 'plt.gca', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (77, 'matplotlib.pyplot.gca', 'plt.gca', ([], {}), True, 'import matplotlib.pyplot as plt\n')]