import os
import os
import pymysql
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
                                                
host = os.environ['MdB_HOST']
user = os.environ['MdB_USER']
passwd = os.environ['MdB_PASSWD']
port = os.environ['MdB_PORT']
                                                
db = pymysql.connect(host=host, user=user, passwd=passwd, port=int(port), db='3MdBs')
cursor = db.cursor()
                                                
result_shock = pd.read_sql("""SELECT shock_params.shck_vel AS shck_vel, 
                                     shock_params.mag_fld AS mag_fld,
                                     log10(emis_VI.NII_6583/emis_VI.HI_6563) AS NII6583_Hb,
                                     log10(emis_VI.OIII_5007/emis_VI.HI_4861) AS OIII5007_Hb
                              FROM shock_params 
                              INNER JOIN emis_VI ON emis_VI.ModelID=shock_params.ModelID
                              INNER JOIN abundances ON abundances.AbundID=shock_params.AbundID
                              WHERE emis_VI.model_type='shock' 
                              AND abundances.name='Allen2008_Solar'
                              AND shock_params.ref='Allen08'
                              AND shock_params.shck_vel BETWEEN 200 AND 1000 
                              AND shock_params.preshck_dens=1
                              ORDER BY shck_vel, mag_fld;""", con=db)
                                                
result_precursor = pd.read_sql("""SELECT shock_params.shck_vel AS shck_vel, 
                                         shock_params.mag_fld AS mag_fld,
                                         log10(emis_VI.NII_6583/emis_VI.HI_6563) AS NII6583_Hb,
                                         log10(emis_VI.OIII_5007/emis_VI.HI_4861) AS OIII5007_Hb
                                  FROM shock_params 
                                  INNER JOIN emis_VI ON emis_VI.ModelID=shock_params.ModelID
                                  INNER JOIN abundances ON abundances.AbundID=shock_params.AbundID
                                  WHERE emis_VI.model_type='shock_plus_precursor' 
                                  AND abundances.name='Allen2008_Solar'
                                  AND shock_params.ref='Allen08'
                                  AND shock_params.shck_vel BETWEEN 200 AND 1000 
                                  AND shock_params.preshck_dens=1
                                 ORDER BY shck_vel, mag_fld;""", con=db)
                                                
fig, ax = plt.subplots(1,2, facecolor='white')
                                                
for B in np.unique(result_shock.mag_fld):
    idx = np.where(result_shock.mag_fld==B)
    ax[0].plot(result_shock.NII6583_Hb[idx[0]], result_shock.OIII5007_Hb[idx[0]], color='black', linewidth=1)
    ax[1].plot(result_precursor.NII6583_Hb[idx[0]], result_precursor.OIII5007_Hb[idx[0]], color='black', linewidth=1)
                                                
for v in np.unique(result_shock.shck_vel):
    idx = np.where(result_shock.shck_vel==v)
    ax[0].plot(result_shock.NII6583_Hb[idx[0]], result_shock.OIII5007_Hb[idx[0]], color='black', linewidth=1)
    ax[1].plot(result_precursor.NII6583_Hb[idx[0]], result_precursor.OIII5007_Hb[idx[0]], color='black', linewidth=1)
                                                
                                                
ax[0].set(xlim=[-0.7, 0.5], 
          ylim=[-0.8, 1.75], 
          xlabel=r'log([N II] $\lambda$6563) / H$\alpha$',
          ylabel=r'log([O III] $\lambda$5007) / H$\beta$')
                                                
ax[1].set(xlim=[-0.7, 0.5], 
          ylim=[-0.8, 1.75], 
          xlabel=r'log([N II] $\lambda$6563) / H$\alpha$')
                                                
ax[0].text(-0.65, 1.6, 'SHOCK ONLY')
ax[1].text(-0.65, 1.6, 'SHOCK + PRECURSOR')
                                                
plt.show()