from discminer.mining_control import _mining_channels_peakint
from discminer.mining_utils import (
    get_noise_mask,
    get_2d_plot_decorators,
    init_data_and_model,
    load_moments,
    load_disc_grid,
    show_output
)

from discminer.cube import Cube
from discminer.model import ReferenceModel
from discminer.plottools import (
    use_discminer_style,
    get_discminer_cmap,
    make_up_ax,
    get_continuous_cmap,
    make_1d_legend
)

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from astropy import units as u

import json
import subprocess

use_discminer_style()

if __name__ == '__main__':
    parser = _mining_channels_peakint(None)
    args = parser.parse_args()

#**************************
#JSON AND SOME DEFINITIONS
#**************************
with open('parfile.json') as json_file:
    pars = json.load(json_file)

meta = pars['metadata']
best = pars['best_fit']
custom = pars['custom']
Rout = best['intensity']['Rout']

nchans = args.nchans
chan_step = args.step

ctitle, clabel, clim, cfmt, cmap_mom, cmap_res, levels_im, levels_cc, unit = get_2d_plot_decorators(args.moment)

#****************
#SOME DEFINITIONS
#****************
file_data = meta['file_data']
tag = meta['tag']
au_to_m = u.au.to('m')

dpc = meta['dpc']*u.pc

mol_tex = {
    '12co': r'$^{12}$CO',
    '13co': r'$^{13}$CO',
    'cs': r'CS'
}

try:
    mol_tex_meta = mol_tex[meta['mol']]
except KeyError:
    mol_tex_meta = meta['mol']

#*******************
#LOAD DATA AND MODEL
#*******************
Rmax = 25*u.au #Larger than sky xmax=18.25au
datacube, model = init_data_and_model(Rmin=0, Rmax=Rmax)
noise_mean, mask = get_noise_mask(datacube, thres=args.sigma)

vchannels = datacube.vchannels
pix_downsamp = model.grid['step']*meta['downsamp_fit']/au_to_m

#Some plot definitions
xmax = model.skygrid['xmax'] 
xlim = np.min([xmax/au_to_m, 25])
extent= np.array([-xmax, xmax, -xmax, xmax])/au_to_m

beam_au = datacube.beam_size.to('au').value
R_lev = np.linspace(args.Rinner, 2.096511, 4)*au_to_m

#**************************
#MAKE MODEL (2D ATTRIBUTES)
#**************************
modelcube = model.make_model(make_convolve=True)

#LOAD DISC GEOMETRY
R, phi, z = load_disc_grid()

Xproj = R[args.surface]*np.cos(phi[args.surface])/au_to_m
Yproj = R[args.surface]*np.sin(phi[args.surface])/au_to_m

r"""#Uncomment for Kelvin
datacube.data /= 1000 
modelcube.data /= 1000
modelcube.convert_to_tb(writefits=False)
datacube.convert_to_tb(writefits=False)
#"""

#*************************
#LOAD MOMENT MAPS
moment_data, moment_model, residuals, mtags = load_moments(args, mask=mask)

if args.moment=='peakintensity': #in Kelvin
    r"""#Uncomment for Kelvin
    moment_data /= 1000
    moment_model /= 1000    
    residuals /= 1000
    #"""
    pass

#*****************
#PLOT CHANNEL MAPS
#*****************
noise_mean, mask = get_noise_mask(datacube, thres=args.sigma) #again, but in Tb units
print ('Mean RMS: %.2e K'%noise_mean)

cmap = None
"""
if 'sum' in meta['kind']:
    import cmasher as cmr
    cmap = cmap_mom = plt.get_cmap('cmr.rainforest_r')
    cmap = plt.get_cmap('cmr.rainforest_r')
"""
idlim = int(0.5*chan_step*(nchans-1))
plot_channels = np.linspace(-idlim,idlim,nchans) + np.argmin(np.abs(vchannels-best['velocity']['vsys']))
plot_channels = np.append([0], plot_channels).astype(int) #dummy channel to make room for peakintensity map later on at ax[0]

kw_channels = dict(channels={'indices': plot_channels}, ncols=nchans+1,
                   xlims = (-xlim, xlim), ylims = (-xlim, xlim), show_beam=False,
                   observable = 'intensity_2',
                   #contours_from=model,
                   #kwargs_contour = dict(colors='w', linewidths=1.0), 
                   unit_coordinates='au', unit_intensity=None, projection=None)
                   
fig, ax = plt.subplots(nrows=8, ncols=nchans+1, figsize=(2.5*(nchans+1), 2.5*8)) #Master figure
plt.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0.05, hspace=0.05)

#CHANNEL MAPS
levels=np.linspace(0*args.sigma*noise_mean, 4*custom['Ilim'], 32)

fig, ax0, im0, cbar0 = datacube.make_channel_maps(fig=fig, ax=ax[0,:], levels=levels, extend='max', annotate_channels=True*0, fmt_cbar='%2d', **kw_channels)
cbar0.set_label('mJy/beam')

temperatures = [25, 100, 150, 200, 400, 600, 800]
subpix = [11, 11, 11, 11, 21, 21, 21]
steps = [20000, 20000, 10000, 10000, 10000, 10000, 10000, 10000]
nT = len(temperatures)

axs = [] #Axes for spectra

nc, nx, ny = datacube.data.shape
nh = int(round(0.5*nx))
nw = 3
dy = 0.02
stat_func = np.sum #i.e. total flux 

data_slice = datacube.data[:, nh-nw:nh+nw+0, nh-nw:nh+nw+0] / datacube.beam_area #convert from mJy/beam to mJy/pixel
data_spec = np.array([stat_func(chan) for chan in data_slice])

data_max = np.max(data_spec)
data_spec/=data_max

nc, sx, sy = data_slice.shape

ax[0, 1].set_ylabel('Folded Data', fontsize=17, color='k')

rectx = model.skygrid['xygrid'][0][nh-nw:nh+nw+0]/au_to_m
xr, yr = rectx[0], rectx[0]
dxr = rectx[-1] - rectx[0]
print ('Rectangle size [au]', dxr)

for i,axi in enumerate(ax[0,1:]):
    rect = patches.Rectangle((xr, yr), dxr, dxr, 
                             facecolor='none', edgecolor='red', 
                             linewidth=1.5, alpha=0.7, fill=False)
    axi.add_patch(rect)
    axi.set_title(r'%.1f$^{\rm km/s}$'%datacube.vchannels[plot_channels[i+1]], fontsize=17)


init_funcs = {
    'velocity': model.velocity_func,
    'z_upper' : model.z_upper_func,
    'z_lower' : model.z_lower_func,
    'intensity' : model.intensity_func,
    'linewidth' : model.linewidth_func,
    'lineslope' : model.lineslope_func,
    'line_profile' : model.line_profile,
    'line_uplow' : model.line_uplow                        
}
    
F_line_arr = []

for i in range(nT):
    stepi = steps[i]
    temp = temperatures[i]
    spix = subpix[i]
    cmd_list = [
        "discminer", "parfile",
        "-f", f"log_pars_hd163296_hcn_0p1_maps_cpd_subpix{spix}_Routmin0p5_PAinit2p5_optthick_ctI0_T{temp}K_cube_60walkers_{stepi}steps.txt",
        "-o", "1"
    ]
    subprocess.run(cmd_list, check=True)
    
    dd, modeli = init_data_and_model(Rmin=0, Rmax=Rmax, subpixels=spix)
    noise_mean, mask = get_noise_mask(datacube, thres=args.sigma)
    modelcube = modeli.make_model(make_convolve=True)

    fig, axi, imi, cbari = modelcube.make_channel_maps(fig=fig, ax=ax[i+1,:], levels=im0[0].levels, annotate_channels=False, cmap=cmap, extend='max', **kw_channels)
    cbari.remove()

    with open('parfile.json') as json_file:
        pars = json.load(json_file)

    best = pars['best_fit']
    Rout = best['intensity']['Rout']
    xc = best['orientation']['xc']
    yc = best['orientation']['yc']    
    vsys_cpd = best['velocity']['vsys']
    
    for axj in ax[i+1, :]:
        modeli.make_disc_axes(axj, Rmax=1.0*Rout*u.au, color='red')
        
    ax[i+1, 1].set_ylabel('T=%dK'%temp, fontsize=17, color='k')

    pos = ax[i+1, -1].get_position()
    axsi = fig.add_axes([pos.x1+0.02, pos.y0+dy, pos.width, pos.height-2*dy])

    model_slice = modelcube.data[:, nh-nw:nh+nw+0, nh-nw:nh+nw+0] / modelcube.beam_area #convert from mJy/beam to mJy/pixel 
    model_spec = np.array([stat_func(chan) for chan in model_slice]) / data_max

    """#Plot all spectra
    for i in range(sx):
        for j in range(sy):
            axsi.step(datacube.vchannels, data_slice[:,i,j], lw=0.1, color='k', where='mid')
            axsi.step(modelcube.vchannels, model_slice[:,i,j], lw=0.2, color='magenta', where='mid')       
    """

    axsi.step(datacube.vchannels, data_spec, lw=2, color='k', where='mid', label='Data')
    axsi.fill_between(datacube.vchannels, data_spec, color='k', step='mid', alpha=0.2)
    
    axsi.step(modelcube.vchannels, model_spec, lw=2, color='goldenrod', where='mid', label='Model')
    axsi.fill_between(modelcube.vchannels, model_spec, color='yellow', step='mid', alpha=0.6)
    
    axsi.axvline(vsys_cpd, ls='--', lw=2.0, dash_capstyle='round', dashes=(2.0, 2.5), color='magenta') #, label=r'$\upsilon_{\rm CPD}$')
     
    axsi.set_ylim(-0.2, 1.2) # 2.5)    
    axsi.tick_params(labelleft=False, left=True, right=True, labelsize=16)        
    axs.append(axsi)

    r"""#Show full resolution model of the CPD velocity field
    #init_header = {}
    modelfullres = ReferenceModel(init_params=best, reset_params=True, init_funcs=init_funcs, Rmin=0, Rmax=Rmax, subpixels=0, npix=200, beam=None)

    vel2d = modelfullres.model.props[0]['upper']
    int2d = modelfullres.model.props[1]['upper'] #masked out beyond Rout
    vel2d = np.where(int2d==0.0, 0, vel2d)

    print (modelfullres.model.extent, extent, model.extent)
    ax[i+1, 0].contourf(vel2d, extent=modelfullres.model.extent, levels=np.linspace(-0.8,0.8,32)+vsys_cpd, cmap=get_discminer_cmap('velocity'))
    ax[i+1, 0].set_xlim(-1.5*Rout+xc, 1.5*Rout+xc)
    ax[i+1, 0].set_ylim(-1.5*Rout+yc, 1.5*Rout+yc)
    #if i==3:
    #    break
    #"""

    #r""" #Compute native fluxes
    modelfullres = ReferenceModel(init_params=best, reset_params=True, init_funcs=init_funcs, Rmin=0, Rmax=Rmax, subpixels=0, npix=200, beam=None)
    I = modelfullres.data
 
    # Pixel sizes (arcsec)
    pix_old = 0.039999999999996004 #I0 is in jy/pixel for this pixel size
    pix_new = modelfullres.pix_size.value*3600 #pixel size of model: 0.0025125628140702
    
    # Scale factor to fix mJy/pixel to the *actual* pixel area
    scale = (pix_new / pix_old)**2   # ~1/253
    
    I_corr = I * scale  # corrected mJy/pixel for high-res pixels
    
    # Now integrate as before
    F_nu = I_corr.sum(axis=(1,2))              # mJy per channel
    dv = np.diff(modelfullres.vchannels)
    dv = np.concatenate(([dv[0]], 0.5*(dv[:-1]+dv[1:]), [dv[-1]]))
    F_line = np.sum(F_nu * dv)                 # mJy km/s
    print(F_line, "mJy km s^-1")

    F_line_arr.append(F_line)
    #"""
    
axs[-1].set_xlabel('l.o.s velocity [km/s]', fontsize=13)
axs[-1].yaxis.set_label_position("right")
axs[-1].set_ylabel('Flux Density', fontsize=13, labelpad=25, rotation=-90)
make_1d_legend(axs[-1], ncol=2, fontsize=15, handlelength=0.8, handletextpad=0.2, columnspacing=0.5)

#DECORATIONS
"""
for axi in ax[:-1,0]:
    axi.set_xlabel('')    
    axi.tick_params(labelleft=False, labelbottom=False)

ax[0,0].set_title('Peak Intensity', color='0.6', fontsize=13, pad=8, loc='center')
ax[0,0].set_ylabel(mol_tex_meta.upper() + ' Data', color='k', fontsize=13, labelpad=6)
ax[1,0].set_ylabel('Model', color='k', fontsize=13, labelpad=4)
"""

#PLOT MOMENT MAP
for axi in ax[:,0]:
    #r"""
    for artist in axi.get_lines() + axi.collections + axi.texts:
        artist.remove()
    #"""
    axi.axis('off')

#kwargs_im = dict(cmap=cmap_mom, extent=extent, levels=im0[0].levels[::1])

#pm0 = ax[0,0].contourf(moment_data, extend='both', **kwargs_im)
#pm1 = ax[1,0].contourf(moment_model, extend='both', **kwargs_im)
#pm2 = ax[2,0].contourf(residuals, cmap=cmap_res, origin='lower', extend='both', extent=extent, levels=np.linspace(-clim, clim, 32))

#PLOT BEAM
"""
for axi in ax[0,:]:
    datacube.plot_beam(axi, fc='gold')        
"""

"""
#for axi in ax[:,int(round(0.5*(nchans+1)))-1]:
for axi in ax[:,0]:
    model.make_emission_surface(
        axi,
        R_lev=R_lev,
        kwargs_R={'colors': '0.1', 'linewidths': 0.5},
        kwargs_phi={'colors': '0.1', 'linewidths': 0.4}
    )
    model.make_disc_axes(axi, Rmax=1.2*Rout*u.au)
    for spine in ['left',  'top', 'right', 'bottom']:
        axi.spines[spine].set_color('0.8')        
    axi.tick_params(which='both', color='0.8')
    #axi.axis('off')
"""

print ('Maximum Flux Density [mJy] from data:', data_max)
print ('temperatures [K] and line fluxes [mJy km s-1]:', temperatures, np.round(F_line_arr, decimals=2))

plt.savefig('channel_maps_peakint_%s_%s_final.png'%(meta['disc'], meta['mol']), bbox_inches = 'tight', dpi=200)    
show_output(args)
