from pathlib import Path
import sys
import numpy
import torch
import torch.nn.functional as F
from PIL import Image, ImageFile
import matplotlib.pyplot as plt
from model import Cnn

sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
import pytorch_gpu_utils


class PredictResult:
    """预测的结果"""

    possibilities: torch.Tensor
    """预测结果，是每个数字不同的概率，是经过softmax后的数值"""

    def __init__(self, possibilities: torch.Tensor):
        self.possibilities = possibilities

    def chosen_number(self) -> int:
        """获取最终选定的数字"""
        # 依然是找最大的那个index
        _, prediction = self.possibilities.max(1)
        return prediction.item()

    def number_possibilities(self) -> list[float]:
        """获取每个数字出现的概率"""
        return list(self.possibilities[0][i].item() for i in range(10))

class Predictor:
    device: torch.device
    model: Cnn

    def __init__(self):
        self.device = pytorch_gpu_utils.get_gpu_device()
        self.model = Cnn().to(self.device)

        # 加载保存好的模型参数
        file_path = Path(__file__).resolve().parent.parent / 'models' / 'cnn.pth'
        self.model.load_state_dict(torch.load(file_path))

    def generic_predict(self, in_data: torch.Tensor) -> PredictResult:
        """
        其它预测函数都要使用的预测后端。其它预测函数将数据处理成Tensor，然后传递给此函数进行实际预测。
        
        :param in_data: 传入的tensor，该tensor的shape必须是28x28，dtype为float32。
        """
        # 上传tensor到GPU
        in_data = in_data.to(self.device)
        # 为了满足要求，要在第一维度挤出2下
        # 一次是灰度通道，一次是批次。
        # 相当于batch size = 1的计算
        in_data = in_data.unsqueeze(0).unsqueeze(0)
        # 开始预测，由于模型输出的是没有softmax的数值，因此最后还需要softmax一下，
        with torch.no_grad():
            out_data = self.model(in_data)
            out_data = F.softmax(out_data, dim=-1)
            return PredictResult(out_data)
       

    def predict_sketchpad(self, image: list[list[bool]]) -> PredictResult:
        input = torch.Tensor(image).float()
        assert(input.dim() == 2)
        assert(input.size(0) == 28)
        assert(input.size(1) == 28)

        return self.generic_predict(input)

    def predict_image(self, image: ImageFile.ImageFile) -> PredictResult:
        # 确保图像为灰度图像，然后转换为numpy数组。
        # 注意这里的numpy数组是只读的，所以要先拷贝一份
        grayscale_image = image.convert('L')
        numpy_data = numpy.reshape(grayscale_image, (28, 28), copy=True)
        # 转换到Tensor，设置dtype
        data = torch.from_numpy(numpy_data).float()
        # 归一化到255，又因为图像输入是白底黑字，需要做转换。
        data.div_(255.0).sub_(1).mul_(-1)

        return self.generic_predict(data)

def main():
    predictor = Predictor()

    # 遍历测试目录中的所有图片，并处理。
    test_dir = Path(__file__).resolve().parent.parent / 'test_images'
    for image_path in test_dir.glob('*.png'):
        if image_path.is_file():
            print(f'Predicting {image_path} ...')
            image = Image.open(image_path)
            rv = predictor.predict_image(image)

            print(f'Predict digit: {rv.chosen_number()}')
            plt.figure(f'Image - {image_path}')
            plt.imshow(image)
            plt.axis('on')
            plt.title(f'Predict digit: {rv.chosen_number()}')
            plt.show()


if __name__ == "__main__":
    pytorch_gpu_utils.print_gpu_availability()
    main()