import gradio as gr
from pyvis.network import Network
import networkx as nx
import numpy as np
import pandas as pd
import os
from datasets import load_dataset
import matplotlib.pyplot as plt

Secret_token = os.getenv('token')

dataset = load_dataset('FDSRashid/hadith_info',data_files = 'Basic_Edge_Information.csv', token = Secret_token, split = 'train')
dataset2 = load_dataset('FDSRashid/hadith_info',data_files = 'Taraf_Info.csv', token = Secret_token, split = 'train')
edge_info = dataset.to_pandas()
taraf_info = dataset2.to_pandas()
cities = taraf_info['City'].unique().tolist()
min_year = int(taraf_info['Year'].min())
max_year = int(taraf_info['Year'].max())
cmap = plt.colormaps['cool']
def value_to_hex(value):
    rgba_color = cmap(value)
    return "#{:02X}{:02X}{:02X}".format(int(rgba_color[0] * 255), int(rgba_color[1] * 255), int(rgba_color[2] * 255))

def subsetEdges(city, year):
  info = taraf_info[(taraf_info['Year'] == year) & (taraf_info['City'] == city)]
  narrators = edge_info[edge_info['Edge_ID'].isin(info['ID'].unique())]
  return narrators
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_visualizer(city, year, num_nodes):
    edges = splitIsnad(subsetEdges(city, year))[['Teacher', 'Student', 'Hadiths']].groupby(['Teacher', 'Student']).sum().reset_index()
    if edges.shape[0] > num_nodes:
        edge_15 = edges.sample(num_nodes)
    else:
        edge_15 = edges.copy()
    teacher_hadiths = edge_15[['Teacher', 'Hadiths']].groupby('Teacher').sum().reset_index()
    net = Network()
    # Create dictionaries to store node roles and colors based on attribute values
    node_roles = {}
    node_colors = {}

    
    for _, row in edge_15.iterrows():
        source = row['Teacher']
        target = row['Student']
        attribute_value = row['Hadiths']
        edge_color = value_to_hex(attribute_value)
        hadith_count = teacher_hadiths[teacher_hadiths['Teacher'] == source]['Hadiths'].to_list()[0]
        edge_color = value_to_hex(attribute_value)
        
        net.add_node(source, color=value_to_hex(hadith_count), font = {'size':30, 'color': 'orange'}, label = f"{source}\nHadiths: {hadith_count}")
        net.add_node(target, color=value_to_hex(attribute_value) , font = {'size': 20, 'color': 'red'}, label = f"{target}\nHadiths: {attribute_value}")
        net.add_edge(source, target, color=edge_color, value=attribute_value)


    net.barnes_hut(gravity=-5000, central_gravity=0.3, spring_length=200)
    html = net.generate_html()
    html = html.replace("'", "\"")

    
    return f"""<iframe style="width: 100%; height: 600px;margin:0 auto" name="result" allow="midi; geolocation; microphone; camera; 
    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
    allow-scripts allow-same-origin allow-popups 
    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
    allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""


with gr.Blocks() as demo:
  Places = gr.Dropdown(choices = cities, value = 'المدينه', label = 'Location')
  Last_Year = gr.Slider(min_year, max_year, value = 9, label = 'End', info = 'Choose the  year to display Narrators')
  num_narrators = gr.Slider(0, 700, value = 400, label = 'Narrators', info = 'Choose the  number of Narrators to display')  
    
  btn = gr.Button('Submit')
  btn.click(fn = network_visualizer, inputs = [Places, Last_Year, num_narrators], outputs = gr.HTML())
  demo.launch()