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

css = """
.table-wrap {
    min-height: 200px;
    max-height: 200px;
}
"""

Secret_token = os.getenv('HF_Token')

edge_info = load_dataset('FDSRashid/hadith_info',token = Secret_token, data_files = 'isnad_info.csv', split = 'train').to_pandas()

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)


def get_node_info(node):
    node_info = narrator_bios[narrator_bios['Rawi ID'] == int(node)]
    # Extract values with defaults, using `.iloc` if available, otherwise default
    student_narrations = node_info['Number of Narrations'].iloc[0] if not node_info.empty else 1
    student_gen = node_info['Generation'].iloc[0] if not node_info.empty else -1
    student_rank = node_info['Narrator Rank'].iloc[0] if not node_info.empty else 'فلان'
    node_name = node_info['Famous Name'].iloc[0] if not node_info.empty else 'فلان'
    return node_info,student_narrations,student_gen, student_rank, node_name

def network_narrator(narrator_id):
    edge_narrator = edge_info[(edge_info['Source'] == str(narrator_id)) | (edge_info['Destination'] == str(narrator_id))]
    edge_full = edge_narrator[['Taraf Count', 'Hadith Count', 'Isnad Count', 'Source', 'Book Count', 'Destination']]
    edge_full['Teacher'] = edge_full['Source'].apply(lambda x: get_node_info(x)[-1] )
    edge_full['Student'] = edge_full['Destination'].apply(lambda x: get_node_info(x)[-1] )
    return edge_full.rename(columns = {'Taraf Count': 'Tarafs', 'Hadith Count':'Hadiths', 'Isnad Count':'Isnads', 'Book Count':'Books'})

def narrator_retriever(name):
    if 'ALL' in name or '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('|'))) | (narrator_bios['Kunya'].apply(lambda x: araby.strip_diacritics(x)).str.contains(araby.strip_diacritics(name), regex=True)) | ((narrator_bios['Laqab'].apply(lambda x: araby.strip_diacritics(x)).str.contains(araby.strip_diacritics(name), regex=True)))]
    

with gr.Blocks(css=css) 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()