#!/usr/bin/env python3
# coding: utf-8

import diskvert as dv
from argparse import ArgumentParser
import numpy as np
import matplotlib.pyplot as plt
import re

parser = ArgumentParser('diskvert-plot',
    description = u'Draw the vertical structure of an accretion disk',
    add_help = False)

parser.add_argument('-xlin', '-x',    action = 'store_true', dest = 'zlin')
parser.add_argument('-xlog', '-xl',   action = 'store_true', dest = 'zlog')
parser.add_argument('-tau',  '-t',    action = 'store_true', dest = 'tau')
parser.add_argument('-col',  '-xcol', action = 'store_true', dest = 'xcol')
parser.add_argument('-grid', action = 'store_true', dest = 'grid')

parser.add_argument('-show', action = 'store_true', dest = 'show')
parser.add_argument('-o', action = 'store', dest = 'outfn')
parser.add_argument('-h', '-help', action = 'help')
parser.add_argument('filename')
args = parser.parse_args()

d,p = dv.col2python(args.filename)

xtype = 'col'
if args.zlin: xtype = 'zlin'
if args.zlog: xtype = 'zlog'
if args.tau: xtype = 'tau'
if args.xcol: xtype = 'col'
if args.grid: xtype = 'grid'

if xtype == 'zlin' or xtype == 'zlog':
    xx = d['z'] / p.zscale
elif xtype == 'col':
    xx = d['nhtot']
elif xtype == 'tau':
    xx = d['tau']
elif xtype == 'grid':
    xx = d['i']
else: raise Exception(u'xtype wrong')

if args.outfn:
    outfn = args.outfn
else:
    outfn = dv.dat2imagefn(args.filename)


hascol = lambda d, kw: kw in d.dtype.names

lis = list()

def setx(ax):
    global lis

    if xtype == 'zlin':
        ax.set_xlabel('$z / H$')
        ax.set_xscale('linear')
        ax.set_xlim(0, (1e3 if p.has_magnetic else 30.))
    elif xtype == 'zlog':
        ax.set_xlabel('$z / H$')
        ax.set_xscale('log')
        ax.set_xlim(0.3, (1e5 if p.has_magnetic else 120.))
    elif xtype == 'grid':
        ax.set_xlabel('$n$')
        ax.set_xscale('linear')
    elif xtype == 'tau':
        ax.set_xlabel('$\\tau$')
        ax.set_xscale('log')
        ax.set_xlim(2e4, 5e-5)
    elif xtype == 'col':
        ax.set_xlabel('$N_H$')
        ax.set_xscale('log')
        ax.set_xlim(1e31, 1e18)
    else: raise Exception(u'xtype wrong')

    if hasattr(p, 'zphot'):
        x = np.interp(p.zphot, d['z'], xx)
        li = ax.axvline(x, linewidth = 0.9, linestyle = '-', color = '#0EBBB1')
        lis.append(li)
    if hasattr(p, 'ztmin'): 
        x = np.interp(p.ztmin, d['z'], xx)
        li = ax.axvline(x,  linewidth = 0.9, linestyle = ':', color = '#cbcf00')
        lis.append(li)
    if hasattr(p, 'ztherm'): 
        x = np.interp(p.ztherm, d['z'], xx)
        li = ax.axvline(x, linewidth = 0.9, linestyle = '--', color = '#F14628')
        lis.append(li)
    # if hasattr(p, 'zflux'): 
    #     x = np.interp(p.zflux, d['z'], xx)
    #     li = ax.axvline(x, linewidth = 0.9, linestyle = ':', color = '#95eb34')
    #     lis.append(li)
    if hasattr(p, 'zinstabil'): 
        x = np.interp(p.zinstabil, d['z'], xx)
        li = ax.axvline(x, linewidth = 0.9, linestyle = ':', color = '#de18ce')
        lis.append(li)

figw = 16.0
fig, axes = plt.subplots(2, 4, figsize = (figw, figw * 0.52), sharex = True, dpi = 96)

tit1 = "$M_{{\\rm BH}}$ = {:.2g}, $\\dot{{m}}$ = {:.2g}, $R / R_{{\\rm schw}}$ = {:.1f}, " \
    .format(p.mbh, p.mdot, p.radius)
tit2 = u'$\\alpha_{{\\rm B}}$ = {:.2g}, $\\eta$ = {:.2g}, $\\nu$ = {:.2g}, $q$ = {:.2f}' \
    .format(p.alpha, p.eta, p.nu, p.qcor) if p.has_magnetic else u'$\\alpha$ = {:.2g}'.format(p.alpha)
plt.suptitle(tit1 + tit2, color = ('#E30E08' if not p.converged else 'black'))

#--------------------------------------------------------------------------#

Mx = lambda x: (x[1:] + x[:-1]) / 2
Dx = lambda x: (x[1:] - x[:-1])

#--------------------------------------------------------------------------#

ax = axes[0,0]
ax.set_title(u'Temperature [K]')

ax.plot(xx, d['trad'], label = '$T_{\\rm rad}$', color = '#B6BBBF')
ax.plot(xx, d['tavg'], label = '$T_{\\rm avg}$', color = '#ABC897', linestyle = '--')
ax.plot(xx, d['temp'], label = '$T$', color = '#DB4024')
ax.plot(Mx(xx), Mx(d['temp']), color = '#aa1111', linewidth = 0.8, linestyle = '-')

setx(ax)

ax.set_yscale('log')
ax.set_ylim(p.teff * 0.5, p.teff * 3e5)
ax.legend(loc = 'best', fontsize = 9)

#--------------------------------------------------------------------------#

ax = axes[0,1]
ax.set_title(u'Pressure [${{\\rm g}} \\ {{\\rm cm}}^{{-1}} \\ {{\\rm s}}^{{-2}}$]')

ax.plot(xx, d['pgas'], label = '$P_{\\rm gas}$', color = '#DB4024')
ax.plot(xx, d['prad'], label = '$P_{\\rm rad}$', color = '#4EBD12')
if p.has_magnetic: 
    ax.plot(xx, d['pmag'], label = '$P_{\\rm mag}$', color = '#2C8BED')
    ax.plot(xx, d['pmagmri'], color = '#8caeed', 
        linestyle = '--', label=r'$P_{\rm mag}^{\rm max}$')

ptot = d['pgas'] + d['prad'] + (d['pmag'] if p.has_magnetic else 0)
ax.plot(xx, ptot, label = '$P_{\\rm tot}$', color = '#1C1C1C')

maxptot = max(ptot)
ax.set_ylim(1e-9 * maxptot, 1.5 * maxptot)
ax.set_yscale('log')
setx(ax)

ax.legend(loc = 'best', fontsize = 9)

ax2 = ax.twinx()
ax2.axhline(1, color = '#aaaaaa', linewidth = 0.8, alpha = 0.7)
ax2.plot(xx, d['radpz'], color = '#888888')
ax2.set_ylim(0, 4)
ax2.set_yscale('linear')


#--------------------------------------------------------------------------#


ax = axes[0,2]
ax.set_title(u'Density [${{\\rm cm}}^{{-3}}$]')

ax.plot(Mx(xx), - 2 * dv.cgs_k_over_mh / 0.5 * Dx(d['rho']) * Mx(d['temp']) / (p.omega**2 * Dx(d['z']**2) * dv.cgs_mhydr),
        color = '#eecccc', linewidth = 1.1, label = '$- (\\Omega^2 z)^{{-1}} T d\\rho / dz$', linestyle='-')

ax.plot(xx, d['dr_pgas'] / dv.cgs_mhydr, color = '#F28D7F', linewidth = 1.2, label = '$- (\\Omega^2 z)^{{-1}} dP_{{\\rm gas}} / dz$')
ax.plot(xx, d['dr_prad'] / dv.cgs_mhydr, color = '#AAEC96', linewidth = 1.2, label = '$- (\\Omega^2 z)^{{-1}} dP_{{\\rm rad}} / dz$')
ax.plot(xx, d['dr_pmag'] / dv.cgs_mhydr, color = '#77B8F2', linewidth = 1.2, label = '$- (\\Omega^2 z)^{{-1}} dP_{{\\rm mag}} / dz$')

ax.plot((xx), (d['rho']) / dv.cgs_mhydr, color = '#1D3608', linewidth = 2.0, label = '$\\rho$', linestyle='-.')
ax.plot(Mx(xx), Mx(d['rho']) / dv.cgs_mhydr, color = '#000000', linewidth = 0.8, linestyle='-')

if 'rho_max' in d.dtype.names:
    ax.plot(xx, d['rho_max'] / dv.cgs_mhydr, color = '#8795a3', linewidth = 1.2, linestyle = '--')

ax.set_yscale('log')
ax.set_ylim(1e-11 * max(d['rho']) / dv.cgs_mhydr, 1.5 * max(d['rho']) / dv.cgs_mhydr)
ax.legend(fontsize = 9, loc = 'best')

setx(ax)

#--------------------------------------------------------------------------#

ax = axes[1,0]
ax.set_title(u'Fluxes')

ax.plot(xx, d['frad'] / p.facc, label = '$F_{\\rm rad}$', color = '#4EBD12')
if p.has_magnetic: ax.plot(xx, d['fmag'] / p.facc, label = '$F_{\\rm mag}$', color = '#2C8BED')
if p.has_conduction: 
    ax.plot(xx,  d['fcnd'] / p.facc, color = '#f0a000', linestyle='-', label = '$F_{\\rm cond}$')
    ax.plot(xx, -d['fcnd'] / p.facc, color = '#f0a000', linestyle='--')
ax.plot(xx, d['heat'] * d['z'] / p.facc, color = '#F069D7', label = '$z \\cdot {{\\cal H}}$')

ax.set_ylim(1e-6, 1.2)
ax.set_yscale('log')


setx(ax)
li1, la1 = ax.get_legend_handles_labels()

if p.has_magnetic:
    ax2 = ax.twinx()
    ax2.axhline(1, color='#cccccc', lw = 1.0)
    ax2.plot(xx, d['qcor'], label = r'$q$', color = '#556655', ls = '--')
    ax2.plot(xx, d['qmri'], label = r'$q_{\rm mri}$', color = '#997799')
    ax2.set_yscale('linear')
    ax2.set_ylim(0,6)
    li2, la2 = ax2.get_legend_handles_labels()
else:
    li2, la2 = [], []

ax.legend(li1 + li2, la1 + la2, fontsize = 9, loc = 'best')

#--------------------------------------------------------------------------#

ax = axes[1,1]
ax.set_title(u'Rates [${{\\rm erg}} \\ {{\\rm cm}}^{{-3}} \\ {{\\rm s}}^{{-1}}$]')

if p.has_conduction:
    ax.plot(Mx(xx), Dx(d['fcnd']) / Dx(d['z']), color='#f0a000', linewidth=1.4, linestyle='-')
    ax.plot(Mx(xx), -Dx(d['fcnd']) / Dx(d['z']), color='#f0a000', linewidth=1.4, linestyle='--')

if hascol(d, 'coolc'):
    ax.plot(xx, d['coolc'], color = '#17CF67', label = '$\\Lambda_C$', linewidth = 1.5)

if p.has_magnetic:
    ax.plot(xx, d['heatr'],  color = '#E03131', label = '${{\\cal H}}_{{\\rm rec}}$', linestyle = '--', linewidth = 1.0)
    ax.plot(xx, d['heatm'],  color = '#2978E2', label = '${{\\cal H}}_{{\\rm mag}}$', linestyle = '--', linewidth = 1.0)

ax.plot(xx, d['cool'],  color = '#444433', label = '${{\\Lambda}}$', linewidth = 2.2, alpha=0.5)
ax.plot(xx, d['heat'],  color = '#212121', label = '${{\\cal H}}$', linewidth = 2.2)

ax.set_yscale('log')

maxheat = max(d['heat'])
ax.set_ylim(maxheat * 1e-7, maxheat * 3)

# if p.has_magnetic:
#     ax2 = ax.twinx()
#     ax2.axhline(1, color='#cccccc', lw = 1.0)
#     ax2.plot(xx, d['qrec'], color = '#997799')
#     ax2.set_yscale('linear')
#     ax2.set_ylim(0,6)

ax.legend(fontsize = 9, loc = 'best')
setx(ax)


#--------------------------------------------------------------------------#

ax = axes[1,2]
# ax.set_title(u'Dimensionless variables')

# ax.axhline(1, ls = ':', c = '#cccccc', lw = 0.8)
# ax.plot(xx, d['taues'], color = '#AAB5B0', label = '$\\tau_{{\\rm es}}$')
# ax.plot(xx, d['tauth'], color = '#CF3912', label = '$\\tau^\\star$')
# if not args.tau: 
#     ax.plot(xx, d['tau'], color = '#042503', linestyle = ':', label = '$\\tau$')
# if p.has_magnetic: 
#     ax.plot(xx, d['betagen'], color = '#0023ad', 
#         linestyle = ':', label = r'$P / P_{\rm mag}$')
# ax.set_yscale('log')
# ax.set_ylim(1e-4, 1e4)

ax.axhline(0, color='#eeeeee')

ax.plot(Mx(xx), Dx(np.log(d['rho'] * d['temp'])) / Dx(np.log(d['z'])), color='#ded445', label='grad rhoT')
ax.plot(Mx(xx), Dx(np.log(d['rho'])) / Dx(np.log(d['z'])), color='#09b84f', label='grad rho')
ax.plot(Mx(xx), Dx(np.log(d['temp'])) / Dx(np.log(d['z'])), color='#ed9879', label='grad T')

A = 1 - d['radpz']
K = dv.cgs_k_over_mh / 0.5

ax.plot(xx, p.omega**2 * d['z']**2 * A / (K * d['temp']), color='#b00505', label='density')
ax.axhline(p.alpha / p.eta, color='#38bce8')


C = p.omega**2 * Mx(d['z'])**2 * (1 - Mx(d['radpz'])) / (K * Mx(d['temp']))
F = (p.alpha / p.eta) * Mx(d['qmri'])
M = p.alpha * p.omega * Mx(d['pgas']) * Mx(d['qmri']) / Mx(d['cool'])
B = (Mx(d['cool_dt']) + M) * (Dx(np.log(d['temp'])) / Dx(np.log(d['z'])) + F + C) \
  - (Mx(d['cool_dr']) + M) * (Dx(np.log(d['rho']))  / Dx(np.log(d['z'])) + F)
D = Mx(d['cool_dt']) - Mx(d['cool_dr'])
kz = B / D
ax.plot(Mx(xx), 1 / kz, color='#dd00ee', linestyle='-', label='k1')

C = p.omega**2 * Mx(d['z'])**2 * (1 - Mx(d['radpz'])) / (K * Mx(d['temp']))
F = 0
M = 0
B = (Mx(d['cool_dt']) + M) * (Dx(np.log(d['temp'])) / Dx(np.log(d['z'])) + F + C) \
  - (Mx(d['cool_dr']) + M) * (Dx(np.log(d['rho']))  / Dx(np.log(d['z'])) + F)
D = Mx(d['cool_dt']) - Mx(d['cool_dr'])
kz = B / D
ax.plot(Mx(xx), 1 / kz, color='#dd00ee', linestyle='--')

# C = p.omega**2 * Mx(d['z'])**2 * (1 - Mx(d['radpz'])) / (K * Mx(d['temp']))
# B = C / (1 - Mx(d['cool_dr']) / Mx(d['cool_dt']))
# D = Dx(np.log(d['rho'] * d['temp'])) / Dx(np.log(d['z']))
# kz = B + D
# ax.plot(Mx(xx), Mx(d['z']) / kz / p.zscale, color='#dd00ee', label='k2', linestyle='--')




ax.set_yscale('symlog', linthresh=1e-2)
ax.set_ylim(-1e4, 1e4)

ax.legend(fontsize = 9, loc = 'best')

# ax2 = ax.twinx()
# compy_color = '#B334D9'
# ax2.plot(xx, d['compy2'], label = u'$y_{{\\rm avg}}$', color = compy_color)
# ax2.set_ylim(0, 0.25)
# ax2.spines['right'].set_color(compy_color)
# ax2.tick_params(axis = 'y', color = compy_color, labelcolor = compy_color)

# lines, labels = ax.get_legend_handles_labels()
# lines2, labels2 = ax2.get_legend_handles_labels()
# ax.legend(lines + lines2, labels + labels2, fontsize = 9, loc = 'best')

setx(ax)

#--------------------------------------------------------------------------#

#--------------------------------------------------------------------------#

ax = axes[1,3]
# ax.plot(xx, d['gradrd'], label = 'rd')
# ax.plot(xx, d['gradad'], label = 'ad')
# ax.set_yscale('symlog', linthresh=1e0)
# ax.legend(fontsize = 9, loc = 'best')

# ax.plot(xx, d['cf_x'] * d['cf_y']**0.5, label = 'ad')
# ax.set_yscale('log')
# ax.set_ylim(1e-2, 1e2)

K = dv.cgs_k_over_mh / 0.5

ax.set_ylim(-1e2, 1e2)
for h in [-1, 0, 1, 2]: ax.axhline(h, color='#eeeeee')


a = - Dx(np.log(d['rho'])) / Dx(np.log(d['temp']))
for i in range(1,a.size):
    if a[i-1] * a[i] < 0: a[i-1:i] = np.nan
ax.plot(Mx(xx), a, color='#0cc5eb', label='grad')


ax.plot(xx, d['cool_dt'] / d['cool_dr'], color='#bcb5e6')

dhd = p.alpha * p.omega * d['pgas'] / d['cool'] #* d['qmri']
ax.plot(xx, (d['cool_dt'] + dhd) / (d['cool_dr'] + dhd), label='max', color='#240cad', linestyle='-.')


# A = 1 - Mx(d['radpz'])
# B = Dx(np.log(d['rho'] * d['temp'])) / Dx(np.log(d['z']))
# C = - p.omega**2 * Mx(d['z'])**2 * A / (K * Mx(d['temp']) * B)
# a = 1 / (1 - C)
# ax.plot(Mx(xx), a, label="gas", color="#c93208")

# A = 1 - Mx(d['radpz'])
# B = Dx(np.log(d['rho'] * d['temp'])) / Dx(np.log(d['z'])) + p.alpha / p.eta
# C = - p.omega**2 * Mx(d['z'])**2 * A / (K * Mx(d['temp']) * B)
# a = 1 / (1 - C)
# ax.plot(Mx(xx), a,  color='#5ce637')


# A = 1 - Mx(d['radpz'])
# B = Dx(np.log(d['rho'] * d['temp'])) / Dx(np.log(d['z'])) + p.alpha / p.eta * Mx(d['qmri'])
# C = - p.omega**2 * Mx(d['z'])**2 * A / (K * Mx(d['temp']) * B)
# a = 1 / (1 - C)
# ax.plot(Mx(xx), a, color='#1d7306', linestyle='--', label='gas+mag')

A = 1 - Mx(d['radpz'])
qq = 0 
B = Dx(np.log(d['rho'])) / Dx(np.log(d['z'])) + p.alpha / p.eta * qq
C = Dx(np.log(d['temp'])) / Dx(np.log(d['z'])) + p.alpha / p.eta * qq \
    + p.omega**2 * Mx(d['z'])**2 * A / (Mx(d['temp']) * K)
a = B / C
ax.plot(Mx(xx), a, color='#debc14', label='gas')

qq = Mx(d['qmri'])
B = Dx(np.log(d['rho'])) / Dx(np.log(d['z'])) + p.alpha / p.eta * qq
C = Dx(np.log(d['temp'])) / Dx(np.log(d['z'])) + p.alpha / p.eta * qq \
    + p.omega**2 * Mx(d['z'])**2 * A / (Mx(d['temp']) * K)
a = B / C
ax.plot(Mx(xx), a, color='#877208', linestyle='--', label='gas+mag')

ahydr = a


# qq = Mx(d['qmri'])
# B = Dx(np.log(d['rho'])) / Dx(np.log(d['z'])) + p.alpha / p.eta * qq \
#     + Dx(np.log(d['temp'])) / Dx(np.log(d['z'])) * 16 * dv.cgs_stef * Mx(d['trad'])**3 / K
# C = Dx(np.log(d['temp'])) / Dx(np.log(d['z'])) + p.alpha / p.eta * qq \
#     + p.omega**2 * Mx(d['z'])**2 * A / (Mx(d['temp']) * K)
# a = B / C
# ax.plot(Mx(xx), a, color='#ff3300', linestyle='--', label='new')






ax.set_yscale('symlog', linthresh = 1e-1)
ax.legend(fontsize=9)
setx(ax)


ax = axes[0,3]
ax.axhline(0, color='#cccccc')
ax.plot(xx, d['cool_dr'], color = '#22bbee', label = 'T = const.', linewidth = 1.2)
ax.plot(xx, d['cool_dt'] , color = '#ed6d5c', label = 'rho = const.', linewidth = 1.5)
# ax.plot(Mx(xx), Mx(d['cool_dt']) - ahydr * Mx(d['cool_dr']), color = '#999999', label = 'mix', linewidth = 2.2)
ax.plot(xx, dhd, color = '#25cf71', label = 'mag', linewidth = 1.4)
ax.plot(xx, d['cool_dt'] - d['cool_dr'], color = '#031340', label = 'P = const.', linewidth = 1.8)


# qq = Mx(d['qmri'])
# C = p.omega**2 * Mx(d['z'])**2 * (1 - Mx(d['radpz'])) / (K * Mx(d['temp']))
# B = Mx(d['cool_dt']) * (Dx(np.log(d['temp'])) / Dx(np.log(d['z'])) + C) - Mx(d['cool_dr']) * Dx(np.log(d['rho'])) / Dx(np.log(d['z']))
# D = Mx(d['cool_dt']) - Mx(d['cool_dr'])
# kk = B / D
# ax.plot(Mx(xx), B, color='#ff9922',  label='k')


ax.set_yscale('symlog', linthresh = 1)
ax.set_ylim(-1e1, 1e7)
ax.legend(fontsize = 9, loc = 'best')
setx(ax)

#--------------------------------------------------------------------------#

# fig.legend([li1, li2, li3], ['$\\tau = 2/3$','$\\tau_{{\\rm cor}}$','$\\tau^\\star = 1$'],
#     fontsize = 9, loc = 'center right')

#--------------------------------------------------------------------------#

# lett = u'abcdefgijklmopqrstuvwxyz'
# for i, ax in enumerate(axes.ravel()):
#     ax.annotate('({})'.format(lett[i]), (0.00, 1.035), xycoords = 'axes fraction', color = 'black', fontsize = 10.5)

#--------------------------------------------------------------------------#

plt.subplots_adjust(0.04, 0.07, 0.98, 0.90, 0.23, 0.21)

plt.savefig(outfn)

if args.show: plt.show()
    
