import sys
import os

sys.path.append(os.path.abspath('..'))
import numpy as np
import random
import cv2
from Utils import HodaDatasetReader


def create_random(a, b):
    return random.randint(a, b)


class PersianStringDigitMaker:
    def __init__(self, n_samples, max_length, target_size: tuple = None):
        self.__train_path = './../Utils/Hoda/DigitDB/Train 60000.cdb'
        self.__test_path = './../Utils/Hoda/DigitDB/Test 20000.cdb'
        self.__extra_path = './../Utils/Hoda/DigitDB/RemainingSamples.cdb'
        self.n_samples = n_samples
        self.max_length = max_length
        if target_size is None:
            self.target_size = (max_length * 32, 32)

    def make(self):
        mean = 0
        std_dev = 1
        num_samples = 1
        res, labels = [], []
        full_data, full_labels = self.__fetch_data()
        for _ in range(self.n_samples):
            length = random.randint(2, self.max_length)
            temp_data, temp_labels = [], []
            for _ in range(length):
                inx = create_random(0, full_data.shape[0])
                digit = full_data[inx]
                label = full_labels[inx]
                lef_right_crop = np.random.normal(mean, std_dev, num_samples)
                x, y = digit.shape
                if lef_right_crop < 0.5:
                    crop = random.randint(0, 10)  # Crop from the right
                    digit = digit[:, :x - crop] if x - crop > 0 else digit
                else:
                    crop = random.randint(0, 10)  # Crop from the left
                    digit = digit[:, crop:] if crop < x else digit
                temp_data.append(
                    cv2.resize(digit, (self.target_size[1], self.target_size[1]), interpolation=cv2.INTER_NEAREST))
                temp_labels.append(label)
            res.append(cv2.resize(np.hstack(temp_data), self.target_size, interpolation=cv2.INTER_NEAREST))
            if len(temp_labels) < self.max_length:
                temp_labels.extend(['-'] * (self.max_length - len(temp_labels)))
            labels.append(temp_labels)
        return np.array(res), np.array(labels)

    def __fetch_data(self):
        train_data, train_labels = HodaDatasetReader.read_hoda_dataset(dataset_path=self.__train_path)
        test_data, test_labels = HodaDatasetReader.read_hoda_dataset(dataset_path=self.__test_path)
        extra_data, extra_labels = HodaDatasetReader.read_hoda_dataset(dataset_path=self.__extra_path)
        full_data = np.concatenate((train_data, test_data, extra_data)).reshape(-1, 32, 32)
        full_labels = np.concatenate((train_labels, test_labels, extra_labels))
        return full_data, full_labels
