import gradio as gr
import pyarabic.araby as araby
import numpy as np
import pandas as pd
import os
from datasets import load_dataset
from datasets import Features
from datasets import Value
from datasets import Dataset
import matplotlib.pyplot as plt



Secret_token = os.getenv('HF_Token')

dataset = load_dataset('FDSRashid/hadith_info',data_files = 'Basic_Edge_Information.csv', token = Secret_token, split = 'train')

lst = ['Rawi ID',
 'Gender',
 'Official Name',
 'Famous Name',
 'Title Name',
 'Kunya',
 'Laqab',
 'Occupation',
 'Wasf',
 'Madhhab',
 'Nasab',
 'Narrator Rank',
 'Generation',
 'Birth Date',
 'Death Date',
 'Age',
 'Place of Stay',
 'Place of Death',
 'Mawla Relation',
 'Famous Relatives',
 'Number of Narrations',
 'Avg. Death Date',
 'Whole Number Death']

dct = {}
for itrm in lst:
  dct[itrm] = Value('string')
dct['Rawi ID'] = Value('int32')
features = Features(dct)
#features = Features({'Rawi ID': Value('int32'), 'Famous Name': Value('string'), 'Narrator Rank': Value('string'), 'Number of Narrations': Value('string'), 'Official Name':Value('string'), 'Title Name':Value('string'), 'Generation': Value('string')} )
narrator_bios = load_dataset("FDSRashid/hadith_info", data_files = 'Teacher_Bios.csv', token = Secret_token,features=features )
narrator_bios = narrator_bios['train'].to_pandas()
narrator_bios.loc[49845, 'Narrator Rank'] = 'رسول الله'
narrator_bios.loc[49845, 'Number of Narrations'] = 0
narrator_bios['Number of Narrations'] = narrator_bios['Number of Narrations'].astype(int)
narrator_bios.loc[49845, 'Number of Narrations'] = 22000 
narrator_bios['Generation'] = narrator_bios['Generation'].replace([None], [-1])
narrator_bios['Generation'] = narrator_bios['Generation'].astype(int)


edge_info = dataset.to_pandas()


def splitIsnad(dataframe):
    teacher_student =dataframe['Edge_Name'].str.split(' TO ')
    dataframe['Teacher'] = teacher_student.apply(lambda x: x[0])
    dataframe['Student'] = teacher_student.apply(lambda x: x[1])
    return dataframe


def network_narrator(narrator_id):
    edge_narrator = edge_info[(edge_info['Teacher_ID'] == narrator_id) | (edge_info['Student_ID'] == narrator_id)]
    edge_full = splitIsnad(edge_narrator[['Tarafs', 'Hadiths', 'Isnads', 'Edge_Name', 'Books']]).drop(['Edge_Name'], axis=1)
    return edge_full

def narrator_retriever(name):
    if 'ALL' in name:
        return narrator_bios
    else:
        full_names = name.replace(', ', '|').replace(',', '|')
        return narrator_bios[(narrator_bios['Official Name'].apply(lambda x: araby.strip_diacritics(x)).str.contains(araby.strip_diacritics(name), regex=True)) | (narrator_bios['Famous Name'].apply(lambda x: araby.strip_diacritics(x)).str.contains(araby.strip_diacritics(name), regex=True)) | (narrator_bios['Rawi ID'].astype(str).isin(full_names.split('|')))]
    

with gr.Blocks() as demo:
    gr.Markdown("Search Narrators using this tool or Retrieve Transmissions involving Narrator")
    with gr.Tab("Search Narrator"):
        text_input = gr.Textbox()
        text_output = gr.DataFrame()
        text_button = gr.Button("Search")
        text_button.click(narrator_retriever, inputs=text_input, outputs=text_output)

    with gr.Tab("View Network"):
        image_input = gr.Number()
        image_button = gr.Button("Retrieve!")
        image_button.click(network_narrator, inputs=[image_input], outputs=[gr.DataFrame(wrap=True)])

    
    

demo.launch()