import concurrent.futures
import glob
import os
import random
import shutil
import subprocess
from datetime import datetime
from os import path, replace, listdir, makedirs, scandir
from os.path import abspath
from pathlib import Path
from time import sleep

import yaml
from ultralytics import YOLO

import constants
from api_client import ApiCredentials, ApiClient
from cdn_manager import CDNCredentials, CDNManager
from dto.annotationClass import AnnotationClass
from inference.onnx_engine import OnnxEngine

from security import Security
from utils import Dotdict
from exports import export_tensorrt, upload_model, export_onnx

today_folder = f'{constants.prefix}{datetime.now():{constants.date_format}}'
train_set = 70
valid_set = 20
test_set = 10
old_images_percentage = 75

DEFAULT_CLASS_NUM = 80
total_files_copied = 0


def form_dataset():
    today_dataset = path.join(constants.config.datasets_dir, today_folder)
    shutil.rmtree(today_dataset, ignore_errors=True)
    makedirs(today_dataset)
    images = []
    with scandir(constants.config.images_dir) as imd:
        for image_file in imd:
            if not image_file.is_file():
                continue
            images.append(image_file)

    print(f'Got {len(images)} images. Start shuffling...')
    random.shuffle(images)

    train_size = int(len(images) * train_set / 100.0)
    valid_size = int(len(images) * valid_set / 100.0)

    print(f'Start copying...')
    copy_annotations(images[:train_size], 'train')
    copy_annotations(images[train_size:train_size + valid_size], 'valid')
    copy_annotations(images[train_size + valid_size:], 'test')


def copy_annotations(images, folder):
    global total_files_copied
    total_files_copied = 0

    def _link_or_copy(src, dst):
        try:
            os.link(src, dst)
        except OSError:
            shutil.copy(src, dst)

    def copy_image(image):
        global total_files_copied
        total_files_copied += 1
        label_name = f'{Path(image.path).stem}.txt'
        label_path = path.join(constants.config.labels_dir, label_name)
        if check_label(label_path):
            _link_or_copy(image.path, path.join(destination_images, image.name))
            _link_or_copy(label_path, path.join(destination_labels, label_name))
        else:
            shutil.copy(image.path, path.join(constants.config.corrupted_images_dir, image.name))
            shutil.copy(label_path, path.join(constants.config.corrupted_labels_dir, label_name))
            print(f'Label {label_path} is corrupted! Copy with its image to the corrupted directory ({constants.config.corrupted_labels_dir})')

        if total_files_copied % 1000 == 0:
            print(f'{total_files_copied} copied...')

    today_dataset = path.join(constants.config.datasets_dir, today_folder)
    destination_images = path.join(today_dataset, folder, 'images')
    makedirs(destination_images, exist_ok=True)

    destination_labels = path.join(today_dataset, folder, 'labels')
    makedirs(destination_labels, exist_ok=True)

    makedirs(constants.config.corrupted_images_dir, exist_ok=True)
    makedirs(constants.config.corrupted_labels_dir, exist_ok=True)

    copied = 0
    print(f'Copying annotations to {destination_images} and {destination_labels} folders:')
    with concurrent.futures.ThreadPoolExecutor() as executor:
        executor.map(copy_image, images)


    print(f'Copied all {copied} annotations to {destination_images} and {destination_labels} folders')


def check_label(label_path):
    if not path.exists(label_path):
        return False
    with open(label_path, 'r') as f:
        lines = f.readlines()
        for line in lines:
            for val in line.split(' ')[1:]:
                if float(val) > 1:
                    return False
    return True


def create_yaml():
    print('creating yaml...')
    lines = ['names:']
    annotation_classes = AnnotationClass.read_json()
    for i in range(DEFAULT_CLASS_NUM):
        if i in annotation_classes:
            lines.append(f'- {annotation_classes[i].name}')
        else:
            lines.append(f'- Class-{i + 1}')
    lines.append(f'nc: {DEFAULT_CLASS_NUM}')

    lines.append(f'test: test/images')
    lines.append(f'train: train/images')
    lines.append(f'val: valid/images')
    lines.append('')

    today_dataset = path.join(constants.config.datasets_dir, today_folder)
    today_yaml = abspath(path.join(today_dataset, 'data.yaml'))
    with open(today_yaml, 'w', encoding='utf-8') as f:
        f.writelines([f'{line}\n' for line in lines])



def resume_training(last_pt_path):
    model = YOLO(last_pt_path)
    results = model.train(data=yaml,                          # path to dataset YAML
                          resume=True,                        # continue from last checkpoint
                          epochs=constants.config.training.epochs,  # total training epochs
                          batch=constants.config.training.batch,    # batch size per GPU
                          imgsz=constants.config.training.imgsz,    # input image size in pixels
                          save_period=constants.config.training.save_period,  # save checkpoint every N epochs
                          workers=constants.config.training.workers,  # dataloader worker threads
                          hsv_h=0.015,                        # hue shift fraction of the color wheel
                          hsv_s=0.7,                          # saturation shift fraction
                          hsv_v=0.4,                          # value (brightness) shift fraction
                          degrees=35.0,                       # max rotation degrees (+/- range)
                          translate=0.1,                      # max translation as fraction of image size
                          scale=0.5,                          # max scale gain (1 +/- scale)
                          shear=10.0,                         # max shear degrees (+/- range)
                          fliplr=0.6,                         # probability of horizontal flip
                          mosaic=1.0,                         # mosaic augmentation probability
                          mixup=0.0)                          # mixup augmentation probability (disabled)


def train_dataset():
    form_dataset()
    create_yaml()
    model = YOLO(constants.config.training.model)

    today_dataset = path.join(constants.config.datasets_dir, today_folder)
    results = model.train(data=abspath(path.join(today_dataset, 'data.yaml')),  # path to dataset YAML
                           epochs=constants.config.training.epochs,             # total training epochs
                           batch=constants.config.training.batch,               # batch size per GPU
                           imgsz=constants.config.training.imgsz,               # input image size in pixels
                           save_period=constants.config.training.save_period,   # save checkpoint every N epochs
                           workers=constants.config.training.workers,           # dataloader worker threads
                           hsv_h=0.015,                                         # hue shift fraction of the color wheel
                           hsv_s=0.7,                                           # saturation shift fraction
                           hsv_v=0.4,                                           # value (brightness) shift fraction
                           degrees=35.0,                                        # max rotation degrees (+/- range)
                           translate=0.1,                                       # max translation as fraction of image size
                           scale=0.5,                                           # max scale gain (1 +/- scale)
                           shear=10.0,                                        # max shear degrees (+/- range)
                           fliplr=0.6,                                          # probability of horizontal flip
                           mosaic=1.0,                                          # mosaic augmentation probability
                           mixup=0.0)                                           # mixup augmentation probability (disabled)

    model_dir = path.join(constants.config.models_dir, today_folder)

    shutil.copytree(results.save_dir, model_dir)
    for file in glob.glob(path.join(model_dir, 'weights', 'epoch*')): # remove unnecessary middle epochs
        os.remove(file)
    shutil.copy(path.join(model_dir, 'weights', 'best.pt'), constants.config.current_pt_model)


def export_current_model():
    export_onnx(constants.config.current_pt_model)
    api_client = ApiClient()
    with open(constants.config.current_onnx_model, 'rb') as binary_file:
        onnx_bytes = binary_file.read()

    key = Security.get_model_encryption_key()
    api_client.upload_big_small_resource(onnx_bytes, 'azaion.onnx', constants.MODELS_FOLDER, key)


if __name__ == '__main__':
    train_dataset()
    export_current_model()
    print('success!')
