
from spectral_cube import SpectralCube, Projection, BooleanArrayMask
import numpy as np
import astropy.units as u
from astropy.io import fits

from astropy.modeling import models, fitting
from pandas import DataFrame
import matplotlib.pyplot as p
import os

from cube_analysis.spectral_stacking import total_profile

from paths import (fourteenB_HI_data_path, fourteenB_HI_file_dict,
                   allfigs_path, alltables_path)

from constants import hi_freq

from plotting_styles import default_figure, onecolumn_figure


hi_cube = SpectralCube.read(fourteenB_HI_file_dict["RotSub_Cube"])
hi_mask = fits.open(fourteenB_HI_file_dict["RotSub_Mask"])[0]
hi_cube = hi_cube.with_mask(hi_mask.data > 0)

hi_cube_cent = SpectralCube.read(fourteenB_HI_file_dict["CentSub_Cube"])
hi_mask_cent = fits.open(fourteenB_HI_file_dict["CentSub_Mask"])[0]
hi_cube_cent = hi_cube_cent.with_mask(hi_mask_cent.data > 0)

hi_cube_peakvel = SpectralCube.read(fourteenB_HI_file_dict["PeakSub_Cube"])
hi_mask_peakvel = fits.open(fourteenB_HI_file_dict["PeakSub_Mask"])[0]
hi_cube_peakvel = hi_cube_cent.with_mask(hi_mask_peakvel.data > 0)

hi_peaktemp_hdu = fits.open(fourteenB_HI_file_dict["PeakTemp"])[0]
hi_peaktemp = Projection.from_hdu(hi_peaktemp_hdu)

hi_beam = hi_cube.beam

dperc = 5
unit = hi_peaktemp.unit
inneredge = np.nanpercentile(hi_peaktemp, np.arange(0, 101, dperc)[:-1]) * unit
outeredge = np.nanpercentile(hi_peaktemp, np.arange(0, 101, dperc)[1:]) * unit
# Add something small to the 100th percentile so it is used
outeredge[-1] += 1e-3 * unit

total_spectrum_hi_peak = \
    np.zeros((inneredge.size, hi_cube.shape[0])) * u.K
total_spectrum_hi_peak_cent = \
    np.zeros((inneredge.size, hi_cube.shape[0])) * u.K
total_spectrum_hi_peak_peakvel = \
    np.zeros((inneredge.size, hi_cube.shape[0])) * u.K

for ctr, (p0, p1) in enumerate(zip(inneredge,
                                   outeredge)):

    print("On bin {} to {} K".format(p0, p1))

    mask = np.logical_and(hi_peaktemp >= p0, hi_peaktemp < p1)

    total_spectrum_hi_peak[ctr] = \
        total_profile(hi_cube, mask, num_cores=4).to(u.K, equivalencies=hi_beam.jtok_equiv(hi_freq))

    total_spectrum_hi_peak_cent[ctr] = \
        total_profile(hi_cube_cent, mask, num_cores=4).to(u.K, equivalencies=hi_beam.jtok_equiv(hi_freq))

    total_spectrum_hi_peak_peakvel[ctr] = \
        total_profile(hi_cube_peakvel, mask, num_cores=4).to(u.K, equivalencies=hi_beam.jtok_equiv(hi_freq))


# We'll make mock SpectralCubes from these so it's easy to calculate
# moments and such from
mask = BooleanArrayMask(np.ones((1178, inneredge.size, 1), dtype=bool),
                        hi_cube.wcs)
rot_stack = SpectralCube(data=total_spectrum_hi_peak.T.reshape((1178, inneredge.size, 1)),
                         wcs=hi_cube.wcs, mask=mask)

cent_stack = SpectralCube(data=total_spectrum_hi_peak_cent.T.reshape((1178, inneredge.size, 1)),
                          wcs=hi_cube.wcs, mask=mask)

peakvel_stack = SpectralCube(data=total_spectrum_hi_peak_peakvel.T.reshape((1178, inneredge.size, 1)),
                             wcs=hi_cube.wcs, mask=mask)

# Now save all of these for future use.
stacked_folder = fourteenB_HI_data_path("stacked_spectra", no_check=True)
if not os.path.exists(stacked_folder):
    os.mkdir(stacked_folder)

wstring = "{}percentile".format(int(dperc))
rot_stack.write(fourteenB_HI_data_path("stacked_spectra/rotation_stacked_peak_{}.fits".format(wstring),
                                       no_check=True), overwrite=True)

cent_stack.write(fourteenB_HI_data_path("stacked_spectra/centroid_stacked_peak_{}.fits".format(wstring),
                                        no_check=True), overwrite=True)