import json
import numpy as np
from PIL import Image
import torch
from torchvision import transforms


def read_json(file_name, suppress_console_info=False):
        with open(file_name, 'r') as f:
            data = json.load(f)
            if not suppress_console_info:
                print("Read from:", file_name)
        return data

def get_file_names(data, imgs_folder, feature_folder, suppress_console_info=False):

    image_file_names = {}
    feature_pathes = {}
    captions = {}
    labels = {}
    lats = {}
    lons = {}

    for img in data['images']:
        image_name = img["image_name"]
        sample_id = img["sample_id"]
        image_id = f'{sample_id}_{image_name}'
        path_data = imgs_folder + f'{sample_id}/{image_name}'
        feature_data = feature_folder + f'{sample_id}/{image_name}.npy'
        # image_file_name.append(path_data)
        # caption.append(img["description"])
        # label.append(img["labels"])
        # lat.append(img["lat"])
        # lon.append(img["lon"])

        image_file_names[image_id] = path_data
        feature_pathes[image_id] = feature_data
        captions[image_id] = img["description"]
        labels[image_id] = img["labels"]
        lats[image_id] = img["lat"]
        lons[image_id] = img["lon"]

    return image_file_names, feature_pathes, captions, labels, lats, lons


def get_data(image_file_names, captions, feature_pathes, labels, lats, lons, image_id):

    image_file_name = image_file_names[image_id]
    feature_path = feature_pathes[image_id]
    caption = captions[image_id]
    label = labels[image_id]
    lat = lats[image_id]
    lon = lons[image_id]
    
    return image_file_name, feature_path, caption, label, lat, lon


def read_by_image_id(data_dir, imgs_folder, feature_folder, image_id=None):
    '''
    return:
        img
        img_ -> transform(img)
        caption
        image_feature -> tensor
        label
        label_en -> text of labels 
        lat
        lon
    '''

    data_info = read_json(data_dir)
    image_file_names, image_features_path, captions, labels, lats, lons = get_file_names(data_info, imgs_folder, feature_folder)

    image_file_name, image_feature_path, caption, label, lat, lon = get_data(image_file_names, captions, image_features_path, labels, lats, lons, image_id)

    label_en = []
    label131 = data_info['labels']

    for lable_name in label131.keys():
        label_id = label131[lable_name]
        for label_singel in label:
            if label_singel == label_id:
                label_en.append(lable_name)
    image_feature = np.load(image_feature_path)

    img = Image.open(image_file_name).convert('RGB')  

    transform = transforms.Compose([
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
        ])

    if transform is not None:
        img_ = np.array(transform(img))
    else:
        img_ = np.array(img)
    img_ = torch.from_numpy(img_.astype('float32'))

    return img, img_, caption, image_feature, label, label_en, lat, lon


# test
data_dir = '/data02/xy/dataEngine/json_data/merged_output_combined_9w_resplit.json'
imgs_folder = '/data02/xy/Clip-hash//datasets/image/'
feature_folder = '/data02/xy/Clip-hash/image_feature/georsclip_21_r0.9_fpn/'
image_id = 'sample44_889.jpg'

# img, img_, caption, image_feature, label, label_en, lat, lon = read_by_image_id(data_dir, imgs_folder, feature_folder, image_id)
# print(img, img_, caption, image_feature, label, label_en, lat, lon)