#!/usr/bin/env python
# -*- coding: utf-8 -*-
# time: 2023/3/20 15:19
# file: figure.py
# author: David
# company: shenyang JY
import sys
import numpy as np
import matplotlib.pyplot as plt


class Figure(object):
    def __init__(self, opt, logger, process):
        self.opt = opt
        self.ds = process
        self.logger = logger

    def get_16_points(self, results):
        # results为模型预测的一维数组，遍历，取每16个点的最后一个点
        preds = []
        for res in results:
            preds.append(res[-1])
        return np.array(preds)

    def draw(self, label_data, predict_norm_data, numbers):
        # label_data = origin_data.data[origin_data.train_num + origin_data.start_num_in_test : ,
        #                                         config.label_in_feature_index]
        # dq_data = dq_data.reshape((-1, self.opt.output_size))
        predict_norm_data = self.get_16_points(predict_norm_data)
        label_data = self.get_16_points(label_data)
        label_data = label_data.reshape((-1, self.opt.output_size))
        # label_data 要进行反归一化
        label_data = label_data * self.ds.std[self.opt.label_in_feature_index] + \
                       self.ds.mean[self.opt.label_in_feature_index]
        predict_data = predict_norm_data * self.ds.std[self.opt.label_in_feature_index] + \
                       self.ds.mean[self.opt.label_in_feature_index]   # 通过保存的均值和方差还原数据
        # dq_data = dq_data * self.ds.std[0] + self.ds.mean[0]
        # predict_data = predict_norm_data
        assert label_data.shape[0] == predict_data.shape[0], "The element number in origin and predicted data is different"

        label_name = [self.ds.tables_column_name[i] for i in self.opt.label_in_feature_index]
        label_column_num = len(self.opt.label_columns)

        # label 和 predict 是错开config.predict_day天的数据的
        # 下面是两种norm后的loss的计算方式，结果是一样的，可以简单手推一下
        # label_norm_data = origin_data.norm_data[origin_data.train_num + origin_data.start_num_in_test:,
        #              config.label_in_feature_index]
        # loss_norm = np.mean((label_norm_data[config.predict_day:] - predict_norm_data[:-config.predict_day]) ** 2, axis=0)
        # logger.info("The mean squared error of stock {} is ".format(label_name) + str(loss_norm))

        loss = np.sum((label_data - predict_data) ** 2)/len(label_data)  # mse
        # loss = np.mean((label_data - predict_data) ** 2, axis=0)
        loss_sqrt = np.sqrt(loss)   # rmse
        loss_norm = 1 - loss_sqrt / self.opt.cap
        # loss_norm = loss/(ds.std[opt.label_in_feature_index] ** 2)
        self.logger.info("The mean squared error of power {} is ".format(label_name) + str(loss_norm))

        # loss1 = np.sum((label_data - dq_data) ** 2) / len(label_data)  # mse
        # loss_sqrt1 = np.sqrt(loss1)  # rmse
        # loss_norm1 = 1 - loss_sqrt1 / self.opt.cap
        # self.logger.info("The mean squared error1 of power {} is ".format(label_name) + str(loss_norm1))
        if self.opt.is_continuous_predict:
            # label_X = range(int((self.ds.data_num - self.ds.train_num - 32)))
            label_X = list(range(numbers))
        else:
            label_X = range(int((self.ds.data_num - self.ds.train_num - self.ds.start_num_in_test)/2))
        print("label_x = ", label_X)
        predict_X = [x for x in label_X]

        if not sys.platform.startswith('linux'):    # 无桌面的Linux下无法输出，如果是有桌面的Linux，如Ubuntu，可去掉这一行
            for i in range(label_column_num):
                plt.figure(i+1)                     # 预测数据绘制
                plt.plot(label_X, label_data[:, i], label='label', color='b')
                plt.plot(predict_X, predict_data[:, i], label='predict', color='g')
                # plt.plot(predict_X, dq_data[:, i], label='dq', color='y')
                # plt.title("Predict actual {} power with {}".format(label_name[i], self.opt.used_frame))
                self.logger.info("The predicted power {} for the last {} point(s) is: ".format(label_name[i], self.opt.predict_points) +
                      str(np.squeeze(predict_data[-self.opt.predict_points:, i])))

                if self.opt.do_figure_save:
                    plt.savefig(self.opt.figure_save_path+"{}predict_{}_with_{}.png".format(self.opt.continue_flag, label_name[i], opt.used_frame))

            plt.show()