# run.py
# Vojtěch Golda

from models import models, colors
import concurrent.futures as cf
from util.generators import dataset_generator, batch_generator, video_frame_generator
from util.running import running
from multiprocessing import get_context
from util.parser import args

import torch
import os
import sys
import cv2 as cv
import signal 
import numpy as np
import time

DEBUG = False
TEST_ALL = False


signal.signal(signal.SIGINT, running.shut_off)

device = torch.device("cpu")
if args.g:
    if not torch.cuda.is_available():
        print("Cuda not available, exiting.", file=sys.stderr)
        quit(1)
    device = torch.device("cuda")

batchsize = args.batch_size
if batchsize < 1:
        print("Invalid batch size. Only positive numbers are accepted.", file=sys.stderr)
        quit(1)
test_accuracy = args.v
inspect_test_data = False
skip_train = False

model_utils = []
if not TEST_ALL:
    ModelUtil = models[args.model_name]
    model_util = ModelUtil(device, args.model_path)
else:
    for m in models:
        model_util = models[m](device)
        model_util.model_model = model_util.get_model(False)
        model_utils.append(model_util)


def train():
    model = model_util.get_model()


    learning_rate = args.learning_rate
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    epoch = 0
    batches = None
    while True if args.epoch_count is None else epoch < args.epoch_count:
        total_loss = 0
        batch_count = 0
        if test_accuracy:
            model.train()
        if not skip_train:
            generator = batch_generator(
                dataset_generator(args.train_dataset), 
                model_util.get_data,
                batchsize
            )
            for batch in generator:
                if not running.running:
                    return
                x, y_batch = batch
                # move to device. Can't paralelize this.
                # handling differs based on heatmap/dict
                if type (y_batch[0]) is dict:
                    x = torch.from_numpy(np.array(x).astype(np.float32)).to(model_util.device)
                    for y in y_batch:
                        y["boxes"] = torch.from_numpy(np.array(y["boxes"]).astype(np.float32)).to(model_util.device)
                        y["scores"] = torch.from_numpy(np.array(y["scores"]).astype(np.float32)).to(model_util.device)
                        y["labels"] = torch.from_numpy(np.array(y["labels"]).astype(np.longlong)).to(model_util.device)
                else:
                    for i in range(len(y_batch)):
                        y_batch[i] = torch.from_numpy(np.array(y_batch[i]).astype(np.float32)).to(model_util.device)
                    for i in range(len(x)):
                        x[i] = torch.from_numpy(np.array(x[i]).astype(np.float32)).to(model_util.device)

                try:
                    loss_dict = model(x, y_batch)
                except torch.cuda.OutOfMemoryError:
                    print("Skipped a batch. Try decreasing the batch size or increasing available memory.", file=sys.stderr)
                    optimizer.zero_grad()
                    continue
                losses = sum(loss for loss in loss_dict.values())
                try:
                    losses.backward()
                    optimizer.step()
                    optimizer.zero_grad()
                except AttributeError:
                    continue
                with torch.no_grad():
                    loss = losses.cpu().numpy()
                    total_loss += loss
                    batch_count += 1
            torch.save(model, model_util.filename)

        if batch_count == 0:
            raise Exception("No usable samples found in " + args.train_dataset)
        epoch_print = f"Epoch {epoch}: average loss {total_loss/batch_count}"
        if epoch >= -1: 
            ##### accuracy
            box_count = 0
            no_box_count = 0
            batch_accuraccy = 0
            if test_accuracy:
                with torch.no_grad():
                    model.eval()
                    for x in batch_generator(
                            dataset_generator(args.validation_dataset), 
                            lambda x:x, 
                            batchsize
                    ):
                        batch_accuraccy+=1
                        try:
                            box_sets = model_util.batch2boxes(model, x)
                        except torch.cuda.OutOfMemoryError:
                            print("Skipped a validation batch. Try decreasing the batch size or increasing available memory.", file=sys.stderr)
                            continue
                        for box_set in box_sets:
                            if len(box_set) == 0:
                                no_box_count += 1
                            box_count+= len(box_set) 
                
                epoch_print += f", validated box count {box_count}, samples with no boxes {no_box_count}" 
        print(epoch_print)
        epoch+=1


def draw_boxes(boxes, image, color, print_boxes=False):

    print_stack = ""
    for b in boxes:
        cv.rectangle(
            image, 
            (b[0], b[1]),
            (b[2], b[3]), 
            color=color, 
            thickness=3
        )

        print_stack += f"{b[0]},{b[1]},{b[2]},{b[3]} "

    if print_boxes:
        print(print_stack)

def display(x, frames, box_sets, wait_for_key, out):
        for o, box_set in enumerate(box_sets):
            for f in frames[o]:
                image = f[0].astype(np.uint8)
                draw_boxes(box_set, image, (0,255,0), args.s)
                if args.w:
                    cv.imshow('Tram detector', image)
                    if wait_for_key:
                        key = cv.waitKey()
                    else:
                        key = cv.waitKey(1)
                    if key == ord('q'):
                        break
                    if key == ord('e'):
                        while cv.waitKey(1) != ord('e'):
                            time.sleep(1)

                    # sleeping with live footage would create a permanent drift
                    if args.camera is not None:
                        time.sleep(1/29)

                if out is not None:
                    out.write(image)

def display_all(x, frames, wait_for_key, out, model_utils, colors):
    # prepare an array of arrays of images to draw the boxes on 
    display_image_sets = []
    for frame_set in frames:
        display_image_set = []
        for frame in frame_set:
            display_image_set.append(frame[0].astype(np.uint8))

        display_image_sets.append(display_image_set)

    # get and draw boxes
    for color_index, model_util in enumerate(model_utils):
        box_sets = model_util.batch2boxes(model_util.model_model, x)
        for o, box_set in enumerate(box_sets):
            for image in display_image_sets[o]:
                draw_boxes(box_set, image, colors[color_index], args.s)

    # display data
    for display_image_set in display_image_sets:
        for image in display_image_set:
            if args.w:
                cv.imshow('Tram detector', image)
                if wait_for_key:
                    key = cv.waitKey()
                else:
                    key = cv.waitKey(1)
                if key == ord('q'):
                    break
                if key == ord('e'):
                    while cv.waitKey(1) != ord('e'):
                        time.sleep(1)

                # sleeping with live footage would create a permanent drift
                if args.camera is not None:
                    time.sleep(1/29)

            if out is not None:
                out.write(image)
            
            if True:
                cv.imwrite(f"out/{running.image_index}.jpg", image)
                running.image_index+=1




def test(input_gen, wait_for_key=False):
    out = None
    running.image_index=0
    if TEST_ALL:
        model_utils_local = model_utils
        colors_local=colors
    else:
        model_utils_local = [model_util] 
        colors_local = [(0, 255, 0)]
        model_util.model_model = model_util.get_model(False)
    if args.output_file is not None:
        if os.system(f"touch {args.output_file} 2> /dev/null") != 0:
            raise Exception(f"Can't create file {args.output_file}. Terminating.")
        
    if args.nth < 1:
        raise ValueError("Invalid n. Only positive numbers are accepted.")

    for x, frames in batch_generator(
            input_gen, 
            lambda x:x, 
            batchsize,
            nth=args.nth
    ):
        if out is None and args.output_file is not None:
            # Define the codec and create VideoWriter object
            # Need to wait for the first frame to get dimensions
            fourcc = cv.VideoWriter_fourcc(*'DIVX')
            out = cv.VideoWriter(args.output_file, fourcc, 29.0, (x[0][0].shape[1],  x[0][0].shape[0]))
        
        display_all(x, frames, wait_for_key, out, model_utils_local, colors_local)
        if not running.running:
            return
    
    #print(running.time_total/63 * 1000)
    running.time_total = 0
    if out is not None:
        out.release()
    cv.destroyAllWindows()


if __name__ == "__main__":
    try:
        if args.t:
            if args.model_name == "TramNet":
                raise Exception("TramNet training is broken. Sorry.")
            try: 
                train()
            except KeyError:
                raise Exception("One or more of the input files are malformed. Terminating.")
        elif args.camera is not None:
            test(video_frame_generator(args.camera))
        elif args.input_file is not None:
            test(video_frame_generator(args.input_file))
        else:
            # test(dataset_generator(args.validation_dataset), True)
            raise Exception("Nothing to do. Terminating. Use --help to display help.")

    except Exception as e:
        if DEBUG:
            raise
        print(e, file=sys.stderr)
        quit(1)
