import sys
import re
import paddle
import subprocess

def check_version():
    try:
        version = paddle.__version__
        version_pattern = r"(\d+)\.(\d+)\.(\d+)"
        match = re.match(version_pattern, version)
        major, minor, patch_version = map(int, match.groups())
        return major, minor
    except ImportError:
        print("PaddlePaddle is not installed.")
        sys.exit(1)

major, minor = check_version()

# PaddlePaddle 2.4及以上版本
if major >= 3 or (major == 2 and minor >= 4):
    from ppdiffusers import StableDiffusionPipeline

    pipe = StableDiffusionPipeline.from_pretrained("./stable-diffusion-2")

    paddle.seed(5232132133)
    prompt = "a portrait of shiba inu with a red cap growing on its head. intricate. lifelike. soft light. sony a 7 r iv 5 5 mm. cinematic post - processing "
    image = pipe(prompt, guidance_scale=7.5, height=768, width=768).images[0]

    image.save("shiba_dog_with_a_red_cap.png")

# PaddlePaddle 2.0到2.3版本
elif 2 <= major < 3 and 0 <= minor < 4:
    import paddle
    import paddle.vision.transforms as T
    from PIL import Image
    import numpy as np
    import json

    model = paddle.vision.models.resnet50(pretrained=True)
    model.eval()

    def preprocess_image(image_path):
        transform = T.Compose([
            T.Resize(256),
            T.CenterCrop(224),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        image = Image.open(image_path).convert('RGB')
        image = transform(image)
        image = paddle.unsqueeze(image, axis=0)
        return image

    def predict(model, image_tensor):
        with paddle.no_grad():
            logits = model(image_tensor)
            probs = paddle.nn.functional.softmax(logits, axis=1).numpy()
        return probs

    def load_labels(filepath):
        with open(filepath, "r") as f:
            labels = json.load(f)
        return labels

    def display_top_predictions(probs, labels):
        top5_idx = np.argsort(probs[0])[-5:][::-1]
        for idx in top5_idx:
            print(f"Class: {labels[idx]}, Probability: {probs[0][idx]:.4f}")

    if __name__ == "__main__":
        image_path = "human.jpg"
        labels_path = "imagenet-simple-labels.json"
        image_tensor = preprocess_image(image_path)
        labels = load_labels(labels_path)
        probs = predict(model, image_tensor)
        display_top_predictions(probs, labels)

# PaddlePaddle 2.0以下版本
else:
    import paddle.fluid as fluid
    import numpy as np
    from PIL import Image

    class SimpleCNN(fluid.dygraph.Layer):
        def __init__(self, num_classes=10):
            super(SimpleCNN, self).__init__()
            self.conv1 = fluid.dygraph.Conv2D(num_channels=3, num_filters=32, filter_size=3, stride=1, padding=1)
            self.pool1 = fluid.dygraph.Pool2D(pool_size=2, pool_stride=2, pool_type='max')
            self.conv2 = fluid.dygraph.Conv2D(num_channels=32, num_filters=64, filter_size=3, stride=1, padding=1)
            self.pool2 = fluid.dygraph.Pool2D(pool_size=2, pool_stride=2, pool_type='max')
            self.fc = fluid.dygraph.Linear(input_dim=64*8*8, output_dim=num_classes)

        def forward(self, x):
            x = self.conv1(x)
            x = fluid.layers.relu(x)
            x = self.pool1(x)
            x = self.conv2(x)
            x = fluid.layers.relu(x)
            x = self.pool2(x)
            x = fluid.layers.flatten(x, 1)
            x = self.fc(x)
            return x

    def load_image(image_path):
        img = Image.open(image_path)
        img = img.resize((32, 32), Image.ANTIALIAS)
        img = np.array(img).astype('float32') / 255.0
        img = img.transpose((2, 0, 1))
        img = np.expand_dims(img, axis=0)
        return img

    def infer(image_path, model):
        with fluid.dygraph.guard():
            img = load_image(image_path)
            img = fluid.dygraph.to_variable(img)
            prediction = model(img)
            predicted_class = np.argmax(prediction.numpy())
            return predicted_class

    if __name__ == "__main__":
        image_path = "human.jpg"
        with fluid.dygraph.guard():
            model = SimpleCNN(num_classes=10)
            model.eval()
            predicted_label = infer(image_path, model)
            print(f"Predicted label: {predicted_label}")
print("finish")