#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Apr 25 13:32:04 2021

@author: vkapoor
"""
from oneat.NEATUtils import plotters
import numpy as np
from oneat.NEATUtils import helpers
from oneat.NEATUtils.helpers import load_json, yoloprediction, normalizeFloatZeroOne, fastnms, averagenms, microscope_dynamic_nms
from keras import callbacks
import os
import tensorflow as tf
import time
from oneat.NEATModels import nets
from oneat.NEATModels.nets import Concat
from oneat.NEATModels.loss import dynamic_yolo_loss
from tqdm import tqdm
# from IPython.display import clear_output
from keras import optimizers
from pathlib import Path
from keras.models import load_model
from tifffile import imread, imwrite, TiffFile, imsave
import csv
from natsort import natsorted
import glob
import matplotlib.pyplot as plt
import h5py
import cv2
import imageio
from PIL import Image
import matplotlib.pyplot as plt
from oneat.NEATModels.neat_goldstandard import NEATDynamic

class NEATPredict(NEATDynamic):
    

    def __init__(self, config, model_dir, model_name, catconfig=None, cordconfig=None):

          super().__init__(config = config, model_dir = model_dir, model_name = model_name, catconfig = catconfig, cordconfig = cordconfig)


    def predict_microscope(self, imagedir, movie_name_list, movie_input, Z_imagedir, Z_movie_name_list, Z_movie_input, start,
                Z_start, downsample=1,
                event_label_interest=1, fileextension='*TIF', nb_prediction=3, n_tiles=(1, 1), Z_n_tiles=(1, 2, 2),
                overlap_percent=0.6, event_threshold=0.5, iou_threshold=0.01, projection_model=None, delay_projection=4,
                fidelity=4, jumpindex = 1):

        self.imagedir = imagedir
        self.basedirResults = self.imagedir + '/' + "live_results"
        Path(self.basedirResults).mkdir(exist_ok=True)
        # Recurrsion variables
        self.movie_name_list = movie_name_list
        self.movie_input = movie_input
        self.Z_movie_name_list = Z_movie_name_list
        self.delay_projection = delay_projection
        self.Z_movie_input = Z_movie_input
        self.Z_imagedir = Z_imagedir
        self.start = start
        self.jumpindex = jumpindex
        self.fidelity = fidelity
        self.event_label_interest = event_label_interest
        self.Z_start = Z_start
        self.projection_model = projection_model
        self.nb_prediction = nb_prediction
        self.fileextension = fileextension
        self.n_tiles = n_tiles
        self.Z_n_tiles = Z_n_tiles
        self.overlap_percent = overlap_percent
        self.iou_threshold = iou_threshold
        self.event_threshold = event_threshold
        self.downsample = downsample
        f = h5py.File(self.model_dir + self.model_name + '.h5', 'r+')
        data_p = f.attrs['training_config']
        data_p = data_p.decode().replace("learning_rate", "lr").encode()
        f.attrs['training_config'] = data_p
        f.close()
        self.model = load_model(self.model_dir + self.model_name + '.h5',
                                custom_objects={'loss': self.yololoss, 'Concat': Concat})

        # Z slice folder listener
        while 1:

            Z_Raw_path = os.path.join(self.Z_imagedir, self.fileextension)
            Z_filesRaw = glob.glob(Z_Raw_path)
            Z_filesRaw = natsorted(Z_filesRaw)

            Raw_path = os.path.join(self.imagedir, '*tif')
            filesRaw = glob.glob(Raw_path)
            filesRaw = natsorted(filesRaw)
            for Z_movie_name in Z_filesRaw:
                Z_Name = os.path.basename(os.path.splitext(Z_movie_name)[0])
                # Check for unique filename
                if Z_Name not in self.Z_movie_name_list and "w2" in Z_Name:
                    self.Z_movie_name_list.append(Z_Name)
                    self.Z_movie_input.append(Z_movie_name)

            for movie_name in filesRaw:
                Name = os.path.basename(os.path.splitext(movie_name)[0])
                # Check for unique filename
                if Name not in self.movie_name_list:
                    self.movie_name_list[Name] = Name
                    self.movie_input[Name] = movie_name

                    total_movies = len(self.movie_input)

            doproject = True

            for i in range(len(self.Z_movie_name_list)):

                Z_Name = self.Z_movie_name_list[i]
                Z_path = self.Z_movie_input[i]

                if Z_Name in self.movie_name_list:

                    Name = self.movie_name_list[Z_Name]
                    path = self.movie_input[Z_Name]

                    doproject = False
                else:
                    doproject = True

                if doproject:
                    time.sleep(self.delay_projection)
                    try:
                        start_time = time.time()
                        print('Reading Z stack for projection')
                        Z_image = imread(Z_path)
                        print('Read properly')
                    except:

                        Z_image = None
                    if Z_image is not None:
                        if self.projection_model is not None:
                            print('Projecting using the projection model')
                            projection = self.projection_model.predict(Z_image, 'ZYX', n_tiles=Z_n_tiles)
                        else:
                            print('Doing max projection')
                            projection = np.amax(Z_image, axis=0)
                        imwrite(self.imagedir + '/' + Z_Name + '.tif', projection.astype('float32'))
                        print(f'____ Projection took {(time.time() - start_time)} seconds ____ ')

                    else:
                        if Z_Name in self.Z_movie_name_list:
                            self.Z_movie_name_list.remove(Z_Name)
                        if Z_movie_name in self.Z_movie_input:
                            self.Z_movie_input.remove(Z_movie_name)

            self.movie_input_list = []
            for (k, v) in self.movie_input.items():
                self.movie_input_list.append(v)
            total_movies = len(self.movie_input_list)
            if total_movies > self.size_tminus + self.start:
                current_movies = imread(self.movie_input_list[self.start:self.start + self.size_tminus + 1])

                sizey = current_movies.shape[1]
                sizex = current_movies.shape[2]
                if self.downsample !=1:
                    scale_percent = 100/self.downsample
                    width = int(sizey * scale_percent / 100)
                    height = int(sizex * scale_percent / 100)
                    dim = (width, height)
                    sizex = height
                    sizey = width

                    current_movies_down = np.zeros([current_movies.shape[0], sizey, sizex])
                    # resize image
                    for j in range(current_movies.shape[0]):
                        current_movies_down[j, :] = cv2.resize(current_movies[j, :], dim, interpolation=cv2.INTER_AREA)
                else:
                    current_movies_down = current_movies
                # print(current_movies_down.shape)
                print('Predicting on Movies:', self.movie_input_list[self.start:self.start + self.size_tminus + 1])
                inputtime = self.start + self.size_tminus

                eventboxes = []
                classedboxes = {}
                smallimage = CreateVolume(current_movies_down, self.size_tminus + 1, 0)

                smallimage = normalizeFloatZeroOne(smallimage, 1, 99.8)
                # Break image into tiles if neccessary
                self.image = smallimage
                print('Doing ONEAT prediction')
                start_time = time.time()
                predictions, allx, ally = self.predict_main(smallimage)
                print(f'____ Prediction took {(time.time() - start_time)} seconds ____ ' )

                # Iterate over tiles
                for p in tqdm(range(0, len(predictions))):

                    sum_time_prediction = predictions[p]

                    if sum_time_prediction is not None:
                        for i in range(0, sum_time_prediction.shape[0]):
                            time_prediction = sum_time_prediction[i]

                            boxprediction = yoloprediction(ally[p], allx[p], time_prediction, self.stride, inputtime,
                                                           self.config, self.key_categories, self.key_cord, self.nboxes,
                                                           'prediction', 'dynamic')

                            if boxprediction is not None:
                                eventboxes = eventboxes + boxprediction

                for (event_name, event_label) in self.key_categories.items():

                    if event_label > 0:
                        current_event_box = []
                        for box in eventboxes:

                            event_prob = box[event_name]
                            if event_prob >= self.event_threshold:
                                current_event_box.append(box)
                        classedboxes[event_name] = [current_event_box]

                self.classedboxes = classedboxes
                self.eventboxes = eventboxes
                print('Performining non maximal supression')
                start_time = time.time()
                self.iou_classedboxes = classedboxes
                self.nms_microscope()
                print(f'____ NMS took {(time.time() - start_time)} seconds ____ ')
                print('Generating ini file')
                self.to_csv_microscope()
                self.predict(self.imagedir, self.movie_name_list, self.movie_input, self.Z_imagedir,
                             self.Z_movie_name_list, self.Z_movie_input, self.start, Z_start,
                             fileextension=self.fileextension, downsample=self.downsample,
                             nb_prediction=self.nb_prediction, n_tiles=self.n_tiles, Z_n_tiles=self.Z_n_tiles,
                             overlap_percent=self.overlap_percent, event_threshold=self.event_threshold,
                             iou_threshold=self.iou_threshold, projection_model=self.projection_model)

    def nms_microscope(self):
        
        
        best_iou_classedboxes = {}
        self.iou_classedboxes = {}
        self.start = self.start + self.jumpindex
        for (event_name,event_label) in self.key_categories.items():
            if event_label > 0:
               
               best_sorted_event_box = microscope_dynamic_nms( self.classedboxes, event_name, self.downsample, self.iou_threshold, self.event_threshold, self.imagex, self.imagey, self.imaget, self.fidelity)
               
               
               best_iou_classedboxes[event_name] = [best_sorted_event_box]
               
        self.iou_classedboxes = best_iou_classedboxes

    def to_csv_microscope(self):

        for (event_name, event_label) in self.key_categories.items():

            if event_label > 0:

                xlocations = []
                ylocations = []
                scores = []
                tlocations = []
                radiuses = []
                confidences = []

                iou_current_event_boxes = self.iou_classedboxes[event_name][0]
                bbox_left_y = self.image.shape[1] // 4
                bbox_right_y = 3 *self.image.shape[1] // 4

                bbox_left_x = self.image.shape[2] // 4
                bbox_right_x = 3 *self.image.shape[2] // 4
                iou_current_event_boxes = sorted(iou_current_event_boxes, key=lambda x: x[event_name], reverse=True)

                for iou_current_event_box in iou_current_event_boxes:
                    xcenter = iou_current_event_box['xcenter']
                    ycenter = iou_current_event_box['ycenter']
                    tcenter = iou_current_event_box['real_time_event']
                    score = iou_current_event_box[event_name]
                    radius = np.sqrt(
                        iou_current_event_box['height'] * iou_current_event_box['height'] + iou_current_event_box[
                            'width'] * iou_current_event_box['width']) // 2
                    confidence = iou_current_event_box['confidence']
                    if xcenter > bbox_left_x and xcenter < bbox_right_x and ycenter > bbox_left_y and ycenter < bbox_right_y:
                        print(round(xcenter), round(ycenter), score)
                        xlocations.append(round(xcenter))
                        ylocations.append(round(ycenter))
                        scores.append(score)
                        tlocations.append(tcenter)
                        radiuses.append(radius)
                        confidences.append(confidence)

                event_count = np.column_stack([xlocations, ylocations])
                total_event_count = np.column_stack([tlocations, ylocations, xlocations, scores, radiuses, confidences])
                csvname = self.basedirResults + "/" + event_name
                if len(xlocations) > 0:
                    writer = csv.writer(open(csvname + ".ini", 'w'))
                    writer.writerow(["[main]"])
                    writer.writerow(["nbPredictions=" + str(self.nb_prediction)])
                    live_event_data = []
                    count = 1

                    for line in event_count:
                        if len(live_event_data) > self.nb_prediction:
                            break
                        live_event_data.append(line)
                        writer.writerow(["[" + str(count) + "]"])
                        writer.writerow(["x=" + str(live_event_data[0][0])])
                        writer.writerow(["y=" + str(live_event_data[0][1])])
                        live_event_data = []

                        count = count + 1

                ImageResults = self.basedirResults + '/' + 'ImageLocations'
                Path(ImageResults).mkdir(exist_ok=True)

                event_data = []
                writer = csv.writer(open(csvname + ".csv", "a"))
                filesize = os.stat(csvname + ".csv").st_size
                if filesize < 1:
                    writer.writerow(['T', 'Y', 'X', 'Score', 'Size', 'Confidence'])
                for line in total_event_count:
                    if line not in event_data:
                        event_data.append(line)
                    writer.writerows(event_data)
                    event_data = []


def CreateVolume(patch, imaget, timepoint):
    starttime = timepoint
    endtime = timepoint + imaget
    smallimg = patch[starttime:endtime, :]

    return smallimg
