import json
import pandas as pd
from collections import defaultdict
from matplotlib import pyplot as plt
import numpy as np

SPLIT_GENRE = {
        "TRAIN": [
            "Casual",
            "Adventure",
            "Action",
            "Indie"
        ],
        "VAL": [
            "Strategy",
            "Education",
            "RPG",
            "Massively Multiplayer",
            "Design & Illustration",
            "Animation & Modeling"
        ],
        "TEST": [
            "Simulation",
            "Sports"
        ]
    }


def parse_img_id(img_id):
    img_id = str(img_id)
    return int(img_id[:-3]), int(img_id[-3:])


def plot(data, img_name, is_genre=False):
    plt.rcParams.update({'font.size': 18})
    colors = []
    for cat in data.keys():
        if cat in SPLIT_GENRE['TRAIN']:
            colors.append('tab:blue')
        elif cat in SPLIT_GENRE['VAL']:
            colors.append('tab:orange')
        elif cat in SPLIT_GENRE['TEST']:
            colors.append('tab:green')
        else:
            colors.append('gray')
    
    
    if is_genre:
        plt.figure(figsize=(10, 6))
        plt.grid(True, axis='y', zorder=0)
        plt.bar(data.keys(), data.values(), color=colors, zorder=3)
        plt.xticks(rotation=30, ha='right')
        plt.legend(handles=[
            plt.Line2D([0], [0], color='tab:blue', lw=10, label='Train'),
            plt.Line2D([0], [0], color='tab:orange', lw=10, label='Val'),
            plt.Line2D([0], [0], color='tab:green', lw=10, label='Test'),
        ])
    else:
        plt.figure(figsize=(15, 6))
        plt.grid(True, axis='y', zorder=0)
        plt.bar(data.keys(), data.values(), zorder=3)
        plt.xticks(rotation=45, ha='right')
    plt.gca().yaxis.set_major_locator(plt.MaxNLocator(integer=True))
    for i, v in enumerate(data.values()):
        plt.text(i, v, str(v), ha='center', va='bottom')
    y_max = max(data.values())
    plt.ylim(0, 1.1 * y_max)
    
    plt.tight_layout()
    plt.savefig(img_name + '.png')
    plt.savefig(img_name + '.pdf')


def plot_genre():

    df = pd.read_csv('app_genre.csv')
    with open('../data/coco_merged/annotations/semantics.json', 'r') as f:
        dataset = json.load(f)

    app_img_map = {}

    for img in dataset['images']:
        img_id = img['id']
        app_id, _ = parse_img_id(img['id'])
        if app_id not in app_img_map:
            app_img_map[app_id] = []
        app_img_map[app_id].append(img_id)

    gnr_app_map = defaultdict(list)

    for i in range(len(df)):
        app_id = df['id'][i]
        tags = df['genre'][i].split(';')
        for tag in tags:
            gnr_app_map[tag].append(int(app_id))


    gnr_app_count = {cat: len(apps) for cat, apps in gnr_app_map.items()}
    gnr_img_count = {cat: sum([len(app_img_map[app]) for app in apps]) for cat, apps in gnr_app_map.items()}
    gnr_anno_count = {cat: sum([len([anno for anno in dataset['annotations'] if anno['image_id'] in app_img_map[app]]) for app in apps]) for cat, apps in gnr_app_map.items()}

    split_genre_order = {genre: i for i, genre in enumerate(SPLIT_GENRE['TRAIN'] + SPLIT_GENRE['VAL'] + SPLIT_GENRE['TEST'])}
    gnr_app_count = {cat: gnr_app_count[cat] for cat in sorted(gnr_app_count, key=lambda x: split_genre_order.get(x, float('inf')))}
    gnr_img_count = {cat: gnr_img_count[cat] for cat in sorted(gnr_img_count, key=lambda x: split_genre_order.get(x, float('inf')))}
    gnr_anno_count = {cat: gnr_anno_count[cat] for cat in sorted(gnr_anno_count, key=lambda x: split_genre_order.get(x, float('inf')))}
    
    
    plot(gnr_app_count, 'genre_app_count', is_genre=True)
    plot(gnr_img_count, 'genre_img_count', is_genre=True)
    plot(gnr_anno_count, 'genre_anno_count', is_genre=True)
    

def plot_cat():

    df = pd.read_csv('app_tag.csv')
    with open('../data/coco_merged/annotations/semantics.json', 'r') as f:
        dataset = json.load(f)

    app_img_map = {}

    for img in dataset['images']:
        img_id = img['id']
        app_id, _ = parse_img_id(img['id'])
        if app_id not in app_img_map:
            app_img_map[app_id] = []
        app_img_map[app_id].append(img_id)

    cat_app_map = defaultdict(list)

    for i in range(len(df)):
        app_id = df['id'][i]
        tags = df['tag'][i].split(';')
        for tag in tags:
            cat_app_map[tag].append(int(app_id))

    cat_app_map.pop('VR')

    cat_app_count = {cat: len(apps) for cat, apps in cat_app_map.items()}
    cat_img_count = {cat: sum([len(app_img_map[app]) for app in apps]) for cat, apps in cat_app_map.items()}
    cat_anno_count = {cat: sum([len([anno for anno in dataset['annotations'] if anno['image_id'] in app_img_map[app]]) for app in apps]) for cat, apps in cat_app_map.items()}

    cat_app_count = {cat: count for cat, count in list(cat_app_count.items())[:30]}
    cat_img_count = {cat: count for cat, count in list(cat_img_count.items())[:30]}
    cat_anno_count = {cat: count for cat, count in list(cat_anno_count.items())[:30]}

    cat_app_count = {cat: count for cat, count in sorted(cat_app_count.items(), key=lambda item: item[1], reverse=True)}
    cat_img_count = {cat: count for cat, count in sorted(cat_img_count.items(), key=lambda item: item[1], reverse=True)}
    cat_anno_count = {cat: count for cat, count in sorted(cat_anno_count.items(), key=lambda item: item[1], reverse=True)}

    plot(cat_app_count, 'tag_app_count')
    plot(cat_img_count, 'tag_img_count')
    plot(cat_anno_count, 'tag_anno_count')


def plot_cat_anno():
    with open('../data/coco_merged/annotations/semantics.json', 'r') as f:
        dataset = json.load(f)
    
    cat_catname_map = {}
    for cat in dataset['categories']:
        cat_catname_map[cat['id']] = cat['name']
        
    cat_anno_map = defaultdict(list)
    for anno in dataset['annotations']:
        cat_anno_map[cat_catname_map[anno['category_id']]].append(anno)
    
    print(len(cat_anno_map['button']))
    cat_anno_map.pop('button')
    
    cat_anno_count = {cat: len(annos) for cat, annos in cat_anno_map.items()}
    cat_anno_count = {cat: count for cat, count in sorted(cat_anno_count.items(), key=lambda item: item[1], reverse=True)}
    cat_anno_count = {cat: count for cat, count in list(cat_anno_count.items())[:30]}
    plot(cat_anno_count, 'ige_cat_anno_count')
    


plot_genre()
plot_cat()
plot_cat_anno()