"""
This module is an example of a barebones QWidget plugin for napari

It implements the Widget specification.
see: https://napari.org/stable/plugins/guides.html?#widgets

Replace code below according to your needs.
"""

import os
import shutil
from pathlib import Path

import napari
import numpy as np
import pandas as pd
import skimage as sm
import tifffile
import timm
import torch
import torch.nn as nn
from fastai.data.block import DataBlock
from fastai.data.transforms import FileSplitter, Normalize, get_image_files
from fastai.layers import Mish
from fastai.optimizer import ranger
from fastai.vision.augment import aug_transforms
from fastai.vision.data import ImageBlock, MaskBlock
from fastai.vision.learner import unet_learner
from fastai.vision.models import resnet101
from fastai.vision.models.unet import DynamicUnet
from magicgui import magic_factory
from napari.layers import Image
from PIL import Image as Image_PIL
from scipy import ndimage
from skimage.feature import blob_log

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# show divisions on focused 2D video
def displayDivisions(filename, DLinput, dfDivisions, orientation=False):
        
    [T, X, Y, rgb] = DLinput.shape
    highlightDivisions = np.zeros([T, 552, 552, 3])

    for x in range(X):
        for y in range(Y):
            highlightDivisions[:, 20 + x, 20 + y, :] = DLinput[:, x, y, :]

    
    for i in range(len(dfDivisions)):
        # for each cell division mark a blue disk around the site
        t, x, y = (
            dfDivisions["T"].iloc[i],
            dfDivisions["X"].iloc[i],
            dfDivisions["Y"].iloc[i],
        )
        x = int(x)
        y = int(y)
        t0 = int(t)

        rr0, cc0 = sm.draw.disk([551 - (y + 20), x + 20], 16)
        rr1, cc1 = sm.draw.disk([551 - (y + 20), x + 20], 11)
        
        # draw line in direction of oriantation of division
        if orientation:
            ori = dfDivisions["Orientation"].iloc[i] * np.pi / 180
            rr2, cc2, val = sm.draw.line_aa(
                int(551 - (y + 16 * np.sin(ori) + 20)),
                int(x + 16 * np.cos(ori) + 20),
                int(551 - (y - 16 * np.sin(ori) + 20)),
                int(x - 16 * np.cos(ori) + 20),
            )

        times = range(t0, t0 + 2)

        timeVid = []
        for t in times:
            if t >= 0 and t <= T - 1:
                timeVid.append(t)

        i = 1
        for t in timeVid:
            highlightDivisions[t][rr0, cc0, 2] = 250 / i
            highlightDivisions[t][rr1, cc1, 2] = 0
            if orientation:
                highlightDivisions[t][rr2, cc2, 2] = 250
            i += 1

    highlightDivisions = highlightDivisions[:, 20:532, 20:532]

    # save file
    highlightDivisions = np.asarray(highlightDivisions, "uint8")
    tifffile.imwrite(
        f"output/divisions{filename}.tif",
        highlightDivisions,
    )


def createFolder(directory):
    try:
        if not os.path.exists(directory):
            os.makedirs(directory)
    except OSError:
        print("Error: Creating directory. " + directory)


def load_checkpoint(checkpoint, model):
    print("=> Loading checkpoint")
    model.load_state_dict(checkpoint["state_dict"])


def sortDL(df, t, x, y):

    a = df[df["T"] == t - 2]
    b = df[df["T"] == t - 1]
    c = df[df["T"] == t + 1]
    d = df[df["T"] == t + 2]

    df = pd.concat([a, b, c, d])

    xMax = x + 13
    xMin = x - 13
    yMax = y + 13
    yMin = y - 13
    if xMax > 511:
        xMax = 511
    if yMax > 511:
        yMax = 511
    if xMin < 0:
        xMin = 0
    if yMin < 0:
        yMin = 0

    dfxmin = df[df["X"] >= xMin]
    dfx = dfxmin[dfxmin["X"] < xMax]

    dfymin = dfx[dfx["Y"] >= yMin]
    df = dfymin[dfymin["Y"] < yMax]

    return df


def intensity(vid, ti, xi, yi):

    [T, X, Y] = vid.shape

    vidBoundary = np.zeros([T, 552, 552])

    for x in range(X):
        for y in range(Y):
            vidBoundary[:, 20 + x, 20 + y] = vid[:, x, y]

    rr, cc = sm.draw.disk([yi + 20, xi + 20], 9)
    div = vidBoundary[ti][rr, cc]
    div = div[div > 0]

    mu = np.mean(div)

    return mu


def maskOrientation(mask):
    S = np.zeros([2, 2])
    X, Y = mask.shape
    x = np.zeros([X, Y])
    y = np.zeros([X, Y])
    x += np.arange(X)
    y += (Y - 1 - np.arange(Y)).reshape(Y, 1)
    A = np.sum(mask)
    Cx = np.sum(x * mask) / A
    Cy = np.sum(y * mask) / A
    xx = (x - Cx) ** 2
    yy = (y - Cy) ** 2
    xy = (x - Cx) * (y - Cy)
    S[0, 0] = -np.sum(yy * mask) / A**2
    S[1, 0] = S[0, 1] = np.sum(xy * mask) / A**2
    S[1, 1] = -np.sum(xx * mask) / A**2
    TrS = S[0, 0] + S[1, 1]
    inertia = np.zeros(shape=(2, 2))
    inertia[0, 0] = 1
    inertia[1, 1] = 1
    q = S - TrS * inertia / 2
    theta = np.arctan2(q[0, 1], q[0, 0]) / 2

    return theta * 180 / np.pi



    # ---------- Find Divisions ----------


@magic_factory(
    Dropdown={
        "choices": [
            "Division heatmap",
            "Division database",
            "Division & orientaton database",
        ]
    }
)
def cellDivision(
    Image_layer: "napari.layers.Image", Dropdown="Division heatmap"
) -> Image:
    filename = Image_layer.name
    resnet = timm.create_model("resnet34", pretrained=True) # Change resnet34 to modify the image classifer which will form the UNetCellDivision model
    resnet.conv1 = nn.Conv2d(
        10, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
    ) # change 10 to modify number of frame inputs

    m = resnet
    m = nn.Sequential(*list(m.children())[:-2])
    model = DynamicUnet(m, 1, (512, 512), norm_type=None).to(DEVICE)
    load_checkpoint(
        torch.load(
            "models/UNetCellDivision10.pth.tar",
            map_location=torch.device("cpu"),
        ),
        model,
    ) # Load trained model 

    DLinput = Image_layer.data
    if len(DLinput.shape) != 4:
        return "ERROR stack must be 4D"
    T, X, Y, rgb = DLinput.shape
    if rgb < 2:
        return """ERROR stack must had two or more colour
         channels (only first 2 will be used)"""

    mask = np.zeros([T, X, Y])
    DLinput = np.array(DLinput)

    for t in range(T - 4): # Aarranges videos into input frames for the deep learning model
        vid = np.zeros([10, X, Y])
        for j in range(5):
            vid[2 * j] = DLinput[int(t + j), :, :, 1]
            vid[2 * j + 1] = DLinput[int(t + j), :, :, 0]

        vid = torch.tensor([vid / 256]).float()
        mask[t + 1] = torch.sigmoid(model(vid)).detach().numpy()

    if Dropdown == "Division heatmap":
        return Image(
            mask,
            name="divisions",
            colormap="blue",
            blending="additive",
            contrast_limits=[0, 0.35],
        )

    img = np.zeros([552, 552])
    img[20:532, 20:532] = mask[0]
    blobs = blob_log(
        img, min_sigma=10, max_sigma=25, num_sigma=25, threshold=30
    )
    blobs_logs = np.concatenate((blobs, np.zeros([len(blobs), 1])), axis=1)

    for t in range(1, T):
        img = np.zeros([552, 552])
        img[20:532, 20:532] = mask[t]
        blobs = blob_log(
            img, min_sigma=10, max_sigma=25, num_sigma=25, threshold=30
        )
        blobs_log = np.concatenate(
            (blobs, np.zeros([len(blobs), 1]) + t), axis=1
        )
        blobs_logs = np.concatenate((blobs_logs, blobs_log))

    label = 0
    _df = []
    for blob in blobs_logs:
        y, x, r, t = blob
        mu = intensity(mask, int(t), int(x - 20), int(y - 20))

        _df.append(
            {
                "Label": label,
                "T": int(t + 1),
                "X": int(x - 20),
                "Y": 532 - int(y),  # map coords without boundary
                "Intensity": mu,
            }
        )
        label += 1

    df = pd.DataFrame(_df)
    createFolder("temp_folder")
    df.to_pickle(f"temp_folder/_dfDivisions{filename}.pkl")
    dfRemove = pd.read_pickle(f"temp_folder/_dfDivisions{filename}.pkl")

    for i in range(len(df)):
        ti, xi, yi = df["T"].iloc[i], df["X"].iloc[i], df["Y"].iloc[i]
        labeli = df["Label"].iloc[i]
        dfmulti = sortDL(df, ti, xi, yi)
        dfmulti = dfmulti.drop_duplicates(subset=["T", "X", "Y"])

        if len(dfmulti) > 0:
            mui = df["Intensity"].iloc[i]
            for j in range(len(dfmulti)):
                labelj = dfmulti["Label"].iloc[j]
                muj = dfmulti["Intensity"].iloc[j]

                if mui < muj:
                    indexNames = dfRemove[dfRemove["Label"] == labeli].index
                    dfRemove.drop(indexNames, inplace=True)
                else:
                    indexNames = dfRemove[dfRemove["Label"] == labelj].index
                    dfRemove.drop(indexNames, inplace=True)

    dfDivisions = dfRemove.drop_duplicates(subset=["T", "X", "Y"])

    if Dropdown == "Division database":
        shutil.rmtree("temp_folder")
        dfDivisions.to_pickle(f"output/dfDivision{filename}.pkl")
        displayDivisions(filename, DLinput, dfDivisions, orientation=False)
        dfDivisions["Y"] = 512 - dfDivisions["Y"] # change coord to match imagej
        dfDivisions.to_excel(f"output/dfDivision{filename}.xlsx")
        return Image(
            mask,
            name="divisions",
            colormap="blue",
            blending="additive",
            contrast_limits=[0, 0.35],
        )

    createFolder("orientationImages")

    for k in range(len(dfDivisions)):
        label = int(dfDivisions["Label"].iloc[k])
        t = int(dfDivisions["T"].iloc[k])
        x = int(dfDivisions["X"].iloc[k])
        y = int(512 - dfDivisions["Y"].iloc[k])

        xMax = int(x + 30)
        xMin = int(x - 30)
        yMax = int(y + 30)
        yMin = int(y - 30)
        if xMax > 512:
            xMaxCrop = 60 - (xMax - 512)
            xMax = 512
        else:
            xMaxCrop = 60
        if xMin < 0:
            xMinCrop = -xMin
            xMin = 0
        else:
            xMinCrop = 0
        if yMax > 512:
            yMaxCrop = 60 - (yMax - 512)
            yMax = 512
        else:
            yMaxCrop = 60
        if yMin < 0:
            yMinCrop = -yMin
            yMin = 0
        else:
            yMinCrop = 0

        vid = np.zeros([10, 120, 120])
        for i in range(5):
            image = np.zeros([60, 60])

            image[yMinCrop:yMaxCrop, xMinCrop:xMaxCrop] = DLinput[
                t - 1 + i, yMin:yMax, xMin:xMax, 1
            ]

            image = np.asarray(image, "uint8")
            tifffile.imwrite("temp_folder/images.tif", image)

            division = Image.open("temp_folder/images.tif")

            division = division.resize((120, 120))
            vid[2 * i] = division

            image = np.zeros([60, 60])

            image[yMinCrop:yMaxCrop, xMinCrop:xMaxCrop] = DLinput[
                t - 1 + i, yMin:yMax, xMin:xMax, 0
            ]

            image = np.asarray(image, "uint8")
            tifffile.imwrite("temp_folder/images.tif", image)

            division = Image.open("temp_folder/images.tif")

            division = division.resize((120, 120))
            vid[2 * i + 1] = division

        vid = np.asarray(vid, "uint8")
        tifffile.imwrite(f"orientationImages/division{label}.tif", vid)

    resnet = timm.create_model("resnet34", pretrained=True) # Change resnet34 to modify the image classifer which will form the UNetCellDivision model
    resnet.conv1 = nn.Conv2d(
        10, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
    ) # change 10 to modify number of frame inputs

    m = resnet
    m = nn.Sequential(*list(m.children())[:-2])
    model = DynamicUnet(m, 1, (120, 120), norm_type=None).to(DEVICE)
    load_checkpoint(
        torch.load(
            "models/UNetOrientation.pth.tar", map_location=torch.device("cpu")
        ),
        model,
    ) # Load trained model 

    _df = []
    for k in range(len(dfDivisions)):
        label = int(dfDivisions["Label"].iloc[k])
        vid = sm.io.imread(f"orientationImages/division{label}.tif").astype(
            int
        )[0]
        vid = torch.tensor([vid / 256]).float()
        mask = torch.sigmoid(model(vid)).detach().numpy()
        ori_mask = maskOrientation(mask)

        _df.append(
            {
                "Label": label,
                "T": dfDivisions["T"].iloc[k],
                "X": dfDivisions["X"].iloc[k],
                "Y": dfDivisions["Y"].iloc[k],
                "Orientation": ori_mask,
            }
        )

    df = pd.DataFrame(_df)

    shutil.rmtree("temp_folder")
    shutil.rmtree("orientationImages")
    df.to_pickle(f"output/dfDivOri{filename}.pkl")
    df["Y"] = 512 - df["Y"] # change coord to match imagej
    df.to_excel(f"output/dfDivOri{filename}.xlsx")
    displayDivisions(filename, DLinput, df, orientation=True)
    return Image(
        mask,
        name="divisions",
        colormap="blue",
        blending="additive",
        contrast_limits=[0, 0.35],
    )