from importlib import import_module
import matplotlib.pyplot as plt
import torch.nn as nn
import numpy as np
import matplotlib
import mlconfig
import torch
import cv2
import gc
import io


def pre_func_resnet18(data=None, src ='', ind_inference=0):
    try:
        matplotlib.use('Agg')
        plt.ioff()

        figsize = (16, 8)
        dpi = 32
        
        # if int(ind_inference) <= 1500:
        #     np.save(src + '_inference_2400_' + str(ind_inference) + '.npy', data)
        
        fig1 = plt.figure(figsize=figsize)
        plt.axes(ylim=(-1, 1))
        sig_real = data[0]
        plt.plot(sig_real, color='black')
        plt.gca().set_axis_off()
        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
        plt.margins(0, 0)
        buf1 = io.BytesIO()
        fig1.savefig(buf1, format="png", dpi=dpi)
        buf1.seek(0)
        img_arr1 = np.frombuffer(buf1.getvalue(), dtype=np.uint8)
        buf1.close()
        img1 = cv2.imdecode(img_arr1, 1)
        img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
        plt.clf()
        plt.cla()
        plt.close()

        fig2 = plt.figure(figsize=figsize)
        plt.axes(ylim=(-1, 1))
        sig_imag = data[1]
        plt.plot(sig_imag, color='black')
        plt.gca().set_axis_off()
        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
        plt.margins(0, 0)
        buf = io.BytesIO()
        fig2.savefig(buf, format="png", dpi=dpi)
        buf.seek(0)
        img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
        buf.close()
        img = cv2.imdecode(img_arr, 1)
        img2 = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        plt.clf()
        plt.cla()
        plt.close()

        img = np.asarray([img1, img2], dtype=np.float32)

        del sig_real
        del sig_imag
        del fig1
        del fig2
        del img1
        del img2
        del buf
        del buf1
        del img_arr
        del img_arr1

        cv2.destroyAllWindows()
        gc.collect()
        print('Подготовка данных завершена')
        print()
        return img

    except Exception as e:
        print(str(e))
        return None


def build_func_resnet18(file_model='', file_config='', num_classes=None):
    try:
        matplotlib.use('Agg')
        plt.ioff()
        torch.cuda.empty_cache()
        config = mlconfig.load(file_config)
        model = getattr(import_module(config.model.architecture.rsplit('.', maxsplit=1)[0]),
                        config.model.architecture.rsplit('.', maxsplit=1)[1])()
        model.conv1 = torch.nn.Sequential(torch.nn.Conv2d(2, 3, kernel_size=(7, 7), stride=(2, 2),
                                                          padding=(3, 3), bias=False), model.conv1)
        model.fc = nn.Linear(in_features=512, out_features=num_classes, bias=True)

        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        if device != 'cpu':
            model = model.to(device)
        model.load_state_dict(torch.load(file_model, map_location=device))
        model.eval()
        cv2.destroyAllWindows()
        gc.collect()
        print('Инициализация модели завершена')
        print()
        return model

    except Exception as exc:
        print(str(exc))
        return None


def inference_func_resnet18(data=None, model=None, mapping=None, shablon=''):
    try:
        matplotlib.use('Agg')
        plt.ioff()
        torch.cuda.empty_cache()
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(device)
        img = torch.unsqueeze(torch.tensor(data), 0).to(device)

        with torch.no_grad():
            output = model(img)
            _, predict = torch.max(output.data, 1)
        prediction = mapping[int(np.asarray(predict.cpu())[0])]
        print('PREDICTION' + shablon + ': ' + prediction)

        output = output.cpu()
        label = np.asarray(np.argmax(output, axis=1))[0]
        output = np.asarray(torch.squeeze(output, 0))
        expon = np.exp(output - np.max(output))
        probability = round((expon / expon.sum())[label], 2)

        del img
        del label
        del expon
        del output
        
        cv2.destroyAllWindows()
        gc.collect()
        print('Уверенность' + shablon + ' в предсказании: ' + str(probability))

        print('Инференс завершен')
        print()
        return [prediction, probability]

    except Exception as exc:
        print(str(exc))
        return None


def post_func_resnet18(src='', model_type='', prediction='', model_id=0, ind_inference=0, data=None):
    try:
        matplotlib.use('Agg')
        plt.ioff()
        '''
        if int(ind_inference) <= 100:
            fig, ax = plt.subplots()
            ax.imshow(data[0], cmap='gray')
            plt.savefig(src + '_inference_' + str(ind_inference) + '_' + prediction + '_real_' + str(model_id) + '_' + model_type + '.png')
            plt.clf()
            plt.cla()
            plt.close(fig)

            del fig
            del ax
            cv2.destroyAllWindows()
            gc.collect()

            fig, ax = plt.subplots()
            ax.imshow(data[1], cmap='gray')
            plt.savefig(src + '_inference_' + str(ind_inference) + '_' + prediction + '_imag_' + str(model_id) + '_' + model_type + '.png')
            plt.clf()
            plt.cla()
            plt.close(fig)

            del fig
            del ax
            cv2.destroyAllWindows()
            gc.collect()
        '''
        plt.clf()
        plt.cla()
        plt.close()
        cv2.destroyAllWindows()
        gc.collect()

        print('Постобработка завершена')
        print()

    except Exception as exc:
        print(str(exc))
        return None