import numpy
import astropy.io.fits
import scipy
import skimage
import pandas
import os
import glob
from numba import jit
import matplotlib.pyplot
import matplotlib.animation
import tqdm 
import multiprocessing
import time
from typing import Union
from pathos.multiprocessing import ProcessingPool

class color:
   PURPLE = '\033[95m'
   CYAN = '\033[96m'
   DARKCYAN = '\033[36m'
   BLUE = '\033[94m'
   GREEN = '\033[92m'
   YELLOW = '\033[93m'
   RED = '\033[91m'
   BOLD = '\033[1m'
   UNDERLINE = '\033[4m'
   END = '\033[0m'


def trackvideo(images_files, cs_files,vmin, vmax, crop=False, filename="video", savepath=""):
    images = [astropy.io.fits.getdata(file, memmap=False) for file in images_files]
    cs = [astropy.io.fits.getdata(file, memmap=False) for file in cs_files]
    if crop:
        images = [image[int(image.shape[0]*0.4):int(image.shape[0]*0.6), int(image.shape[1]*0.4):int(image.shape[1]*0.6)] for image in images]
        cs = [c[int(c.shape[0]*0.4):int(c.shape[0]*0.6), int(c.shape[1]*0.4):int(c.shape[1]*0.6)] for c in cs]
    ims = []
    fig, ax= matplotlib.pyplot.subplots(1, 1, figsize=(10, 5))

    for j in range(len(cs)):
        im1 = ax.imshow(images[j], origin='lower', cmap='gray', vmin=vmin, vmax=vmax)
        cs1 = ax.contour(cs[j], colors="red", origin='lower')

        ims.append([im1, cs1])


    ani = matplotlib.animation.ArtistAnimation(fig, ims, interval=50, blit=True,
                                    repeat_delay=1000)
    ani.save(savepath+f'{filename}.mp4', writer='ffmpeg', fps=5)


def extrusion(feature, df, asc_files, return_center=False):
    n = feature
    idx = df.index[df['label'] == n].tolist()[0]
    # get the corresponding frames from the dataframe
    frames = df.Frames[idx]
    x_center = df['X'][idx]
    y_center = df['Y'][idx]
    max_y = numpy.max(numpy.diff(x_center))
    max_x = numpy.max(numpy.diff(y_center))
    max_disp = int(numpy.max([max_x, max_y]))

    imgs = [astropy.io.fits.getdata(asc_files[i], memmap=False) for i in range(len(frames))]

    masks = [im == n for im in imgs]
    cropped_masks = [mask[int(y_center[0]-max_disp):int(y_center[0]+max_disp), int(x_center[0]-max_disp):int(x_center[0]+max_disp)] for mask in masks]
    x_center = numpy.array(x_center)
    y_center = numpy.array(y_center)
    stacked_masks = numpy.stack(cropped_masks, axis=2)
    z_center = numpy.array(frames)
    if return_center:
        x_center = x_center - (x_center[0]-max_disp)
        y_center = y_center - (y_center[0]-max_disp)
        return stacked_masks, x_center, y_center, z_center
    else:
        return stacked_masks, z_center
    

def follow_video(stacked_masks, savepath="", filename="follow_video"):
    ims = []
    fig, ax= matplotlib.pyplot.subplots(1, 1, figsize=(10, 5))
    for j in range(numpy.shape(stacked_masks)[2]):
        im1 = ax.imshow(stacked_masks[:,:,j], origin='lower', cmap='gray', vmin=0, vmax=1)
        ims.append([im1])

    ani = matplotlib.animation.ArtistAnimation(fig, ims, interval=50, blit=True,
                                    repeat_delay=1000)
    ani.save(savepath+f'{filename}.mp4', writer='ffmpeg', fps=5)


def freq_map(datapath, minlifetime=15, flag="x", title=None):
    df = pandas.read_json(datapath+"dataframe.json")
    df = df[df["Lifetime"] > minlifetime]
    df = df[df["stdVx"] < 10]
    df = df[df["stdVy"] < 10]
    df.reset_index(drop=True, inplace=True)
    img = astropy.io.fits.getdata(datapath+"00-data/0000.fits")
    cs = astropy.io.fits.getdata(datapath+"03-assoc/0000.fits")
    temp_img = numpy.copy(cs)
    lifetime = numpy.array(df["Lifetime"])

    peaks = []
    x = []
    y = []
    area = []


    for j in range(len(df)):
        x_ = numpy.array(df["X"].iloc[j]).mean()
        y_ = numpy.array(df["Y"].iloc[j]).mean()
        area_ = numpy.array(df["Area"].iloc[j]).mean()
        x.append(x_)
        y.append(y_)
        area.append(area_)
        if flag=="x":
            f, Pxx = scipy.signal.periodogram(numpy.array(df["Vx"].iloc[j]), fs=1/45, nfft=256, detrend="linear", scaling="density", )
            Pxx = Pxx/numpy.max(Pxx)
            peaks.append(f[numpy.argmax(Pxx)])
        elif flag=="y":
            f, Pyy = scipy.signal.periodogram(numpy.array(df["Vy"].iloc[j]), fs=1/45, nfft=256, detrend="linear", scaling="density", )
            Pyy = Pyy/numpy.max(Pyy)
            peaks.append(f[numpy.argmax(Pyy)])


    diameter = numpy.array(numpy.sqrt(numpy.array(area)/numpy.pi)*2).astype(int)
    peaks = numpy.array(peaks)

    peaks_img = numpy.zeros_like(temp_img).astype(numpy.float64)

    for i in range(len(peaks)):
        x_ = int(x[i])
        y_ = int(y[i])
        peaks_img[y_, x_] = peaks[i]
        for j in range(-diameter[i]//2, diameter[i]//2):
            for k in range(-diameter[i]//2, diameter[i]//2):
                if numpy.sqrt(j**2 + k**2) < diameter[i]//2:
                    peaks_img[y_+j, x_+k] = peaks[i]



    # smooth the peaks_img so that its all colored 
    from scipy.ndimage import gaussian_filter
    peaks_img = gaussian_filter(peaks_img, sigma=1)

    matplotlib.pyplot.figure(figsize=(10, 10))
    matplotlib.pyplot.imshow(img, cmap="gray", vmin=-100, vmax=100)
    # replace the values of the peaks_img with 0 to nan to make the background transparent
    peaks_img[peaks_img < 0.0007] = numpy.nan
    peaks_img = numpy.array(peaks_img)*1000

    matplotlib.pyplot.imshow(peaks_img, cmap="jet", alpha=0.99)
    matplotlib.pyplot.colorbar(label="Frequency (mHz)")
    if title != None:
        matplotlib.pyplot.savefig(f"{title}.png")

    

def amp_map(datapath, minlifetime=15, flag="x", title=None):
    df = pandas.read_json(datapath+"dataframe.json")
    df = df[df["Lifetime"] > minlifetime]
    df = df[df["stdVx"] < 10]
    df = df[df["stdVy"] < 10]
    df.reset_index(drop=True, inplace=True)
    img = astropy.io.fits.getdata(datapath+"00-data/0000.fits")
    cs = astropy.io.fits.getdata(datapath+"03-assoc/0000.fits")
    temp_img = numpy.copy(cs)
    lifetime = numpy.array(df["Lifetime"])

    peaks = []
    x = []
    y = []
    area = []


    for j in range(len(df)):
        x_ = numpy.array(df["X"].iloc[j]).mean()
        y_ = numpy.array(df["Y"].iloc[j]).mean()
        vx_ = numpy.array(df["Vx"].iloc[j]).std()
        vy_ = numpy.array(df["Vy"].iloc[j]).std()
        area_ = numpy.array(df["Area"].iloc[j]).mean()
        x.append(x_)
        y.append(y_)
        area.append(area_)
        if flag=="x":
            peaks.append(vx_)
        elif flag=="y":
            peaks.append(vy_)


    diameter = numpy.array(numpy.sqrt(numpy.array(area)/numpy.pi)*2).astype(int)
    peaks = numpy.array(peaks)
    # keep only values below 10


    peaks_img = numpy.zeros_like(temp_img).astype(numpy.float64)

    for i in range(len(peaks)):
        x_ = int(x[i])
        y_ = int(y[i])
        peaks_img[y_, x_] = peaks[i]
        for j in range(-diameter[i]//2, diameter[i]//2):
            for k in range(-diameter[i]//2, diameter[i]//2):
                if numpy.sqrt(j**2 + k**2) < diameter[i]//2:
                    peaks_img[y_+j, x_+k] = peaks[i]



    # smooth the peaks_img so that its all colored 
    from scipy.ndimage import gaussian_filter

    peaks_img = gaussian_filter(peaks_img, sigma=1)
    peaks_img[peaks_img < 0.001] = numpy.nan


    matplotlib.pyplot.figure(figsize=(10, 10))
    matplotlib.pyplot.imshow(img, cmap="gray", vmin=-100, vmax=100)
    # replace the values of the peaks_img with 0 to nan to make the background transparent


    matplotlib.pyplot.imshow(peaks_img, cmap="jet", alpha=0.99)
    matplotlib.pyplot.colorbar(label="Amplitude (km/s)")
    if title != None:
        matplotlib.pyplot.savefig(f"{title}.png")
