import os
import tempfile
import shutil

import dlc2kinematics as dlc
import gradio as gr
import matplotlib.cm as cm
import pandas as pd
from dlc2kinematics import compute_speed, load_data
from matplotlib import pyplot as plt


def all_plot(bodyparts, speed_csv, tmpdir):

   # speed_csvからグラフを作成。
   plt.figure()
   for i in bodyparts:
      plt.plot(speed_csv[i], label=i, color=cm.rainbow(int(255 * bodyparts.index(i) / len(bodyparts))))

   plt.legend()
   plt.xlabel('frame')
   plt.ylabel('speed(px/frame)')

   # 横軸の最大値を取得して設定
   xmax = speed_csv.index.max()
   plt.xlim(0, xmax)
   plt.ylim(0, 100)

   # グラフの凡例をグラフ外に表示。
   plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, fontsize=10)

   # 凡例をすべて画像内に収める
   plt.tight_layout()
   plt.savefig(f'{tmpdir}/all_plot.png')


def bodyparts_plot(bodyparts, speed_csv, tmpdir):
   # speed_csvからグラフを作成。
   for i in bodyparts:
      plt.figure()
      plt.plot(speed_csv[i], label=i, color=cm.rainbow(int(255 * bodyparts.index(i) / len(bodyparts))))
      plt.legend()
      plt.xlabel('frame')
      plt.ylabel('speed')

      # 横軸の最大値を取得して設定
      xmax = speed_csv.index.max()
      plt.xlim(0, xmax)
      plt.ylim(0, 100)

      # グラフの凡例をグラフ外に表示。
      plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, fontsize=10)
      # 凡例がはみ出ないように
      plt.tight_layout()

      plt.savefig(f'{tmpdir}/{i}.png')

def likelihood_plot(bodyparts, likelihood_csv, tmpdir):
   for i in bodyparts:
      plt.figure()
      plt.plot(likelihood_csv[i], label=i, color=cm.rainbow(int(255 * bodyparts.index(i) / len(bodyparts))))
      plt.legend()
      plt.xlabel('frame')
      plt.ylabel('likeloohood')
  
      # 横軸の最大値を取得して設定
      xmax = likelihood_csv.index.max()
      plt.xlim(0, xmax)
      plt.ylim(0, 1)
  
      # グラフの凡例をグラフ外に表示。
      plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, fontsize=10)
  
      # 凡例をすべて画像内に収める
      plt.tight_layout()
  
      plt.savefig(f'{tmpdir}/{i}.png')

def all_likelihood_plot(bodyparts, likelihood_csv, tmpdir):
   plt.figure()
   for i in bodyparts:
      plt.plot(likelihood_csv[i], label=i, color=cm.rainbow(int(255 * bodyparts.index(i) / len(bodyparts))))

   plt.legend()
   plt.xlabel('frame')
   plt.ylabel('likeloohood')

   # 横軸の最大値を取得して設定
   xmax = likelihood_csv.index.max()
   plt.xlim(0, xmax)
   plt.ylim(0, 1)

   # グラフの凡例をグラフ外に表示。
   plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, fontsize=10)

   # 凡例をすべて画像内に収める
   plt.tight_layout()

   plt.savefig(f'{tmpdir}/likelihood_plot.png')


def main(h5_file):
   h5_file_name = os.path.basename(h5_file.name)
   h5_file_name = os.path.splitext(h5_file_name)[0]
   df, bd, scorer = load_data(h5_file.name)
   speed = compute_speed(df, ['all'])

   # speedのカラムを取得
   new_columns = speed.columns.droplevel([0, 2])
   speed.columns = new_columns
   # カラムを取得
   bodyparts = df.columns.droplevel([0, 2]).to_list()
   bodyparts = list(dict.fromkeys(bodyparts))
   speed_csv = pd.DataFrame()
   likelihood_csv = pd.DataFrame()
   for i in bodyparts:
      speed_csv[i] = speed[i].iloc[:, 0]
      likelihood_csv[i] = speed[i].iloc[:, 1]

   with tempfile.TemporaryDirectory(dir=".") as tmpdir:
      bodyparts_plot(bodyparts, speed_csv, tmpdir)
      all_plot(bodyparts, speed_csv, tmpdir)
      all_likelihood_plot(bodyparts, likelihood_csv, tmpdir)
      likelihood_plot(bodyparts, likelihood_csv, tmpdir)
      shutil.make_archive(f"{h5_file_name}", 'zip', root_dir=tmpdir)
      return f"{h5_file_name}.zip"


iface = gr.Interface(fn=main, inputs="file", outputs="file")
iface.launch()