import pandas as pd
import numpy as np
import gradio as gr
import unicodedata
import os

REGISTERS_DIR = 'registers'
SEASONS = [os.path.splitext(file)[0] for file in os.listdir(REGISTERS_DIR)]

def load_register(season):
  register = pd.read_csv(f'registers/{season}.csv', index_col='jp_name')
  register.en_name = register.en_name.str.replace(',', '')
  return register

def translate_name(jp_name, register):
  jp_name = unicodedata.normalize('NFKC', jp_name)
  en_name = np.nan
  
  if jp_name in register.index:
    en_name = register.loc[jp_name].en_name
  elif jp_name in register.index.map(lambda name: name.split('.')[-1]):
    tmp_register = register.copy()
    tmp_register.index = tmp_register.index.map(lambda name: name.split('.')[-1])
    en_name = tmp_register.loc[jp_name].en_name

  if isinstance(en_name, pd.DataFrame):
    unique_names = en_name.en_name.unique()
    en_name = unique_names[0]
  elif not isinstance(en_name, str):
    en_name = np.nan

  return en_name

def translate_names(season, jp_names):
  register = load_register(season)
  return '\n'.join([str(translate_name(jp_name, register)) for jp_name in jp_names.splitlines()])

demo = gr.Interface(fn=translate_names, inputs=[gr.Dropdown(choices=SEASONS), "text"], outputs=gr.Text(show_copy_button=True))
demo.launch(debug=True)