# -*- coding: UTF-8 -*-
import os
import platform
import pandas as pd
import matplotlib

from src.common.path import ordered_list_json_dir
from src.common.json import load_json, dump_json
from src.common.gif import concat_gif_list, merge_video_audio
from src.bar_chart_race.chart import bar_chart_race
from dataclasses import dataclass


@dataclass
class InputMeta:
    '''
    type: 指定类型，如果是 "json_str" 表示一个JSON文件夹，如果是"csv"表示一个csv文件
    JSON 文件夹:
        约定每个文件的文件名是月份，每个JSON文件是一个数组，数组元素是标签统计信息
        name_field: 指定标签名字的字段名
        count_field: 指定标签月份统计信息的字段名
    CSV 文件
        month_field: 指定月份字段
        name_field: 指定标签名字的字段名
        count_field: 指定标签月份统计信息的字段名
        audio: 音频
    '''
    type: str
    path: str
    month_field: str
    name_field: str
    count_field: str
    audio: str


@dataclass
class OutputMeta:
    '''
    输出配置
    path: 输出路径
    title: 标题
    x_label: X轴名字
    y_label: Y轴名字
    month_count: 绘制月份，用来调试，使用较少的月份快速查看输出效果
    '''
    path: str
    ext: str
    title: str
    x_label: str
    y_label: str
    month_count: int


class Top:
    '''
    # TopN 竞赛动图
    '''

    def __init__(self, input: InputMeta, output: OutputMeta):
        self.input_type = input.type
        self.input = input.path
        self.month_field = input.month_field
        self.name_field = input.name_field
        self.count_field = input.count_field

        if os.path.exists(input.audio):
            self.audio = input.audio
        else:
            self.audio = None

        self.output = output.path
        self.ext = output.ext
        self.title = output.title
        self.x_label = output.x_label
        self.y_label = output.y_label
        self.month_count = output.month_count

        self.override_all_gif = False

    def build(self):
        self.__load_dataframe()
        # self.__exit()
        self.__config_font()

        max_rows = self.df.shape[0]
        i = 0
        j = 0
        df = self.df
        gifs = []
        os.makedirs(self.output, exist_ok=True)
        while i < max_rows:
            end = i+12
            if end >= max_rows:
                end = max_rows+1
            step = end-i

            filename = os.path.join(self.output, f'{j}.{self.ext}')
            if i+step >= max_rows:
                # 最后一个
                last_df = df[i:end]

                # 生成一个短摘要
                min_half = 5
                if min_half > last_df.shape[0]:
                    min_half = 0
                self.df = last_df[min_half:]
                filename_abstracts = os.path.join(
                    self.output, f'{j}_abstracts.{self.ext}')
                self.__build_race(filename_abstracts)

                # 加强最后一帧
                self.df = last_df
                for k in range(0, 12):
                    self.df = self.df.append(df[end-2:end])
                self.__build_race(filename)
            else:
                self.df = df[i:end]
                self.__build_race(filename)

            gifs.append(filename)
            i += step
            j += 1

        # 合并 gif 生成mp4
        all = f'{self.output}.mp4'
        if os.path.exists(all):
            ret = input(f"文件:{all}已存在，是否覆盖？[y/n]:")
            if ret == 'y':
                concat_gif_list(gifs, all)
        else:
            concat_gif_list(gifs, all)

        # 合成 mp4 和 音频
        if self.audio is not None:
            all_with_audio = f'{self.output}_with_audio.mp4'
            if os.path.exists(all_with_audio):
                ret = input(f"文件:{all_with_audio}已存在，是否覆盖？[y/n]:")
                if ret == 'y':
                    merge_video_audio(all, self.audio, all_with_audio)
            else:
                merge_video_audio(all, self.audio, all_with_audio)
        else:
            print(f"音频文件{self.audio}不存在，请自行准备一个合适的音频文件。")

    def __load_dataframe(self):
        if self.input_type == 'json_dir':
            self.__load_dataframe_from_json()
        elif self.input_type == 'csv':
            self.__load_dataframe_from_csv()
        else:
            raise Exception(
                "un suport input type, only support json_dir or csv")

    def __load_dataframe_from_json(self):
        tag_files = ordered_list_json_dir(self.input)
        tag_array = []
        month_count = 0
        for tag_file in tag_files:
            tags = load_json(tag_file)
            tag_month = os.path.basename(tag_file)[0:10]
            print(tag_month)
            for tag in tags:
                tag_array.append([
                    tag_month,
                    tag[self.name_field],
                    tag[self.count_field]
                ])
            month_count += 1
            if self.month_count is not None:
                if month_count == self.month_count:
                    break

        df = pd.DataFrame(tag_array, columns=[
            'month',
            self.name_field,
            self.count_field]
        )

        self.df = df.pivot_table(
            index='month',
            columns=self.name_field,
            values=self.count_field,
            fill_value=0
        )
        print(self.df.head())

    def __load_dataframe_from_csv(self):
        df = pd.read_csv(
            self.input,
            encoding='utf_8_sig',
            index_col=[0],
            lineterminator="\n"
        )

        max_rows = df.shape[0]
        if self.month_count is not None:
            max_rows = min(max_rows, self.month_count)
        df = df[0:max_rows]

        self.df = df.pivot_table(
            index=self.month_field,
            columns=self.name_field,
            values=self.count_field,
            fill_value=0
        )

        print(self.df.head())

    def __config_font(self):
        # matplotlib.rc("font", family='FZFangSong-Z02', weight='bold')
        if(platform.system() == 'Linux'):
            matplotlib.font_manager.findfont('WenQuanYi Micro Hei')
            matplotlib.rc("font", family='WenQuanYi Micro Hei')
        else:
            matplotlib.rc("font", family='FZFangSong-Z02')

    def __build_race(self, filename):
        build = False
        if self.override_all_gif:
            build = True
        else:
            if os.path.exists(filename):
                ret = input(f"文件:{filename}已存在，是否覆盖？[a/y/n]:")
                if ret == 'y':
                    build = True
                elif ret == 'a':
                    build = True
                    self.override_all_gif = True
            else:
                build = True

        if not build:
            return

        bar_chart_race(
            df=self.df,
            filename=filename,
            n_bars=10,
            title=self.title,
            title_size=24,
            logo=self.x_label,
            y_label=self.y_label,
            # filter_column_colors=True,
            figsize=(8, 8),
            bar_size=0.7,
            bar_label_size=14,
            tick_label_size=14,
            steps_per_period=20,
        )

    def __exit(self):
        import sys
        sys.exit()
