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()