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, 3, 6, 42), 'pandas.read_csv', 'pd.read_csv', ({(6, 15, 6, 41): '"""sevn_output/output_0.csv"""'}, {}), "('sevn_output/output_0.csv')", True, 'import pandas as pd\n'), ((13, 4, 13, 54), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n'), ((39, 8, 39, 52), 'pandas.concat', 'pd.concat', ({(39, 18, 39, 51): '[dt[idx0].Mass_0, dt[idx1].Mass_1]'}, {}), '([dt[idx0].Mass_0, dt[idx1].Mass_1])', True, 'import pandas as pd\n'), ((40, 10, 40, 56), 'pandas.concat', 'pd.concat', ({(40, 20, 40, 55): '[dt[idxb0].Mass_0, dt[idxb1].Mass_1]'}, {}), '([dt[idxb0].Mass_0, dt[idxb1].Mass_1])', True, 'import pandas as pd\n'), ((41, 12, 41, 58), 'pandas.concat', 'pd.concat', ({(41, 22, 41, 57): '[dt[idxm0].Mass_0, dt[idxm1].Mass_1]'}, {}), '([dt[idxm0].Mass_0, dt[idxm1].Mass_1])', True, 'import pandas as pd\n'), ((44, 12, 44, 58), 'pandas.concat', 'pd.concat', ({(44, 22, 44, 57): '[dt[idx0].Mzams_0, dt[idx1].Mzams_1]'}, {}), '([dt[idx0].Mzams_0, dt[idx1].Mzams_1])', True, 'import pandas as pd\n'), ((45, 14, 45, 62), 'pandas.concat', 'pd.concat', ({(45, 24, 45, 61): '[dt[idxb0].Mzams_0, dt[idxb1].Mzams_1]'}, {}), '([dt[idxb0].Mzams_0, dt[idxb1].Mzams_1])', True, 'import pandas as pd\n'), ((46, 16, 46, 64), 'pandas.concat', 'pd.concat', ({(46, 26, 46, 63): '[dt[idxm0].Mzams_0, dt[idxm1].Mzams_1]'}, {}), '([dt[idxm0].Mzams_0, dt[idxm1].Mzams_1])', True, 'import pandas as pd\n'), ((49, 9, 49, 43), 'pandas.concat', 'pd.concat', ({(49, 19, 49, 42): '[dt[idx0].a, dt[idx1].a]'}, {}), '([dt[idx0].a, dt[idx1].a])', True, 'import pandas as pd\n'), ((50, 11, 50, 47), 'pandas.concat', 'pd.concat', ({(50, 21, 50, 46): '[dt[idxb0].a, dt[idxb1].a]'}, {}), '([dt[idxb0].a, dt[idxb1].a])', True, 'import pandas as pd\n'), ((51, 13, 51, 49), 'pandas.concat', 'pd.concat', ({(51, 23, 51, 48): '[dt[idxm0].a, dt[idxm1].a]'}, {}), '([dt[idxm0].a, dt[idxm1].a])', True, 'import pandas as pd\n'), ((55, 0, 55, 26), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((57, 0, 57, 18), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(57, 12, 57, 13): '(1)', (57, 14, 57, 15): '(2)', (57, 16, 57, 17): '(1)'}, {}), '(1, 2, 1)', True, 'import matplotlib.pyplot as plt\n'), ((58, 0, 58, 68), 'matplotlib.pyplot.scatter', 'plt.scatter', (), '', True, 'import matplotlib.pyplot as plt\n'), ((59, 0, 59, 75), 'matplotlib.pyplot.scatter', 'plt.scatter', (), '', True, 'import matplotlib.pyplot as plt\n'), ((60, 0, 60, 81), 'matplotlib.pyplot.scatter', 'plt.scatter', (), '', True, 'import matplotlib.pyplot as plt\n'), ((62, 0, 62, 17), 'matplotlib.pyplot.xscale', 'plt.xscale', ({(62, 11, 62, 16): '"""log"""'}, {}), "('log')", True, 'import matplotlib.pyplot as plt\n'), ((63, 0, 63, 17), 'matplotlib.pyplot.yscale', 'plt.yscale', ({(63, 11, 63, 16): '"""log"""'}, {}), "('log')", True, 'import matplotlib.pyplot as plt\n'), ((64, 0, 64, 45), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (), '', True, 'import matplotlib.pyplot as plt\n'), ((65, 0, 65, 55), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (), '', True, 'import matplotlib.pyplot as plt\n'), ((67, 0, 67, 23), 'matplotlib.pyplot.legend', 'plt.legend', (), '', True, 'import matplotlib.pyplot as plt\n'), ((69, 0, 69, 18), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(69, 12, 69, 13): '(1)', (69, 14, 69, 15): '(2)', (69, 16, 69, 17): '(2)'}, {}), '(1, 2, 2)', True, 'import matplotlib.pyplot as plt\n'), ((70, 0, 70, 65), 'matplotlib.pyplot.scatter', 'plt.scatter', (), '', True, 'import matplotlib.pyplot as plt\n'), ((71, 0, 71, 71), 'matplotlib.pyplot.scatter', 'plt.scatter', (), '', True, 'import matplotlib.pyplot as plt\n'), ((72, 0, 72, 77), 'matplotlib.pyplot.scatter', 'plt.scatter', (), '', True, 'import matplotlib.pyplot as plt\n'), ((73, 0, 73, 17), 'matplotlib.pyplot.xscale', 'plt.xscale', ({(73, 11, 73, 16): '"""log"""'}, {}), "('log')", True, 'import matplotlib.pyplot as plt\n'), ((74, 0, 74, 17), 'matplotlib.pyplot.yscale', 'plt.yscale', ({(74, 11, 74, 16): '"""log"""'}, {}), "('log')", True, 'import matplotlib.pyplot as plt\n'), ((75, 0, 75, 56), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (), '', True, 'import matplotlib.pyplot as plt\n'), ((76, 0, 76, 45), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (), '', True, 'import matplotlib.pyplot as plt\n'), ((81, 0, 81, 18), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((82, 0, 82, 28), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(82, 12, 82, 27): '"""analysis3.png"""'}, {}), "('analysis3.png')", True, 'import matplotlib.pyplot as plt\n'), ((83, 0, 83, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((61, 9, 61, 27), 'numpy.linspace', 'np.linspace', ({(61, 21, 61, 22): '(0)', (61, 23, 61, 26): '(140)'}, {}), '(0, 140)', True, 'import numpy as np\n'), ((61, 28, 61, 46), 'numpy.linspace', 'np.linspace', ({(61, 40, 61, 41): '(0)', (61, 42, 61, 45): '(140)'}, {}), '(0, 140)', True, 'import numpy as np\n'), ((66, 0, 66, 9), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((77, 0, 77, 9), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n')]