import re
from collections import Counter, defaultdict
import gradio as gr
import pandas as pd
import plotly.graph_objects as go
import networkx as nx
import plotly.express as px

def create_recommendation_system(pairs):
    sets_by_length = defaultdict(list)

    for pair in pairs:
        sets_by_length[len(pair)].append(set(pair))

    sorted_keys = sorted(sets_by_length.keys(), reverse=True)

    def recommend(restaurant):
        recommended = set()
        for length in sorted_keys:
            for rest_set in sets_by_length[length]:
                if restaurant in rest_set:
                    recommended.update(rest_set - {restaurant})
        return list(recommended)

    return recommend

def recommend(input_string):
    f = open("task1_output_n_c1_sup4.txt", "r")
    data = f.read()
    data = data.split("Frequent Itemsets:")

    clean_string = re.sub(r'\n', '', data[1])
    pattern = r"\([^\)]+\)"

    extracted_tuples = re.findall(pattern, clean_string.replace("'", ""))

    tuple_list = [tuple(s.strip('()').split(', ')) for s in extracted_tuples]

    recommendation_system = create_recommendation_system(tuple_list)

    output_list = recommendation_system(input_string)

    images_labels = [("https://logowik.com/content/uploads/images/aw-restaurants5299.jpg", "A&W Restaurant"),
    ("https://inspirebrands.com/wp-content/uploads/2017/10/Arbys.jpg","Arbys"),
    ("https://www.shutterstock.com/image-vector/vinnytsia-ukraine-october-8-2023-600nw-2372329457.jpg","Burger King"),
    ("https://static.wixstatic.com/media/cc8696_2f0a91e1d9d245efa98b100adfcdd887~mv2.png/v1/crop/x_0,y_0,w_242,h_71/fill/w_339,h_94,al_c,lg_1,q_85,enc_auto/logo.png","California Kitchen"),
    ("https://logos-world.net/wp-content/uploads/2022/11/Carls-Jr.-Logo-500x281.png","Carls Jr."),
    ("https://upload.wikimedia.org/wikipedia/commons/thumb/0/02/Chick-fil-A_Logo.svg/873px-Chick-fil-A_Logo.svg.png","Chick-Fil-A"),
    ("https://upload.wikimedia.org/wikipedia/en/thumb/3/38/CECLogo2019.svg/330px-CECLogo2019.svg.png","Chuck E. Cheese"),
    ("https://upload.wikimedia.org/wikipedia/commons/thumb/c/cc/Culver%27s_logo.svg/558px-Culver%27s_logo.svg.png","Culvers"),
    ("https://upload.wikimedia.org/wikipedia/commons/thumb/a/ae/Dairy_Queen_logo.svg/300px-Dairy_Queen_logo.svg.png", "Dairy Queen"),
    ("https://logowik.com/content/uploads/images/493_dominospizza.jpg", "Dominos Pizza"),
    ("https://logos-world.net/wp-content/uploads/2023/03/Five-Guys-Logo-500x281.png", "Five Guys"),
    ("https://entrackr.com/storage/2023/06/Good-Flippin.jpg", "Good Flippin Burger"),
    ("https://logowik.com/content/uploads/images/hardees4024.jpg", "Hardees"),
    ("https://logos-marques.com/wp-content/uploads/2023/04/In-N-Out-Burger-Logo-768x432.png", "In-N-Out"),
    ("https://logos-world.net/wp-content/uploads/2022/08/Jack-in-the-Box-Logo-500x281.png", "Jack in the Box"),
    ("https://1000logos.net/wp-content/uploads/2021/05/Jollibee-logo-768x432.png", "Jollibee"),
    ("https://logowik.com/content/uploads/images/674_kfc.jpg", "KFC"),
    ("https://visitguernseycounty.com/wp-content/uploads/2022/05/Little-Caesars-Logo-1024x576.jpg", "Little Caesars"),
    ("https://logowik.com/content/uploads/images/mcdonalds-icon.jpg", "McDonalds"),
    ("https://images.squarespace-cdn.com/content/v1/53a47e51e4b0e78ae9ed2e97/4b8df8ab-662a-4381-92a6-97839099d5a7/Papa-Johns-logo.jpg", "Papa Johns"),
    ("https://cdn.nwe.io/files/x/d5/cc/c9810c9ad13188a3e0b2bc6577d7.jpg", "Pizza Inn"),
    ("https://logowik.com/content/uploads/images/294_pizza_hut_new_logo.jpg", "Pizza hut"),
    ("https://upload.wikimedia.org/wikipedia/commons/a/a6/PizzaExpress_Logo.jpg", "PizzaExpress"),
    ("https://images.squarespace-cdn.com/content/v1/53a47e51e4b0e78ae9ed2e97/be0aea01-84cc-4b89-a5e8-314ed394dc3c/Popeyes+logo.jpg", "Popeyes"),
    ("https://upload.wikimedia.org/wikipedia/commons/d/d9/Original_Round_Table_Pizza_Logo.jpg", "Round Table Pizza"),
    ("https://logowik.com/content/uploads/images/328_sbarro.jpg", "Sbarro"),
    ("https://1000logos.net/wp-content/uploads/2023/04/Shake-Shack-Logo-768x432.png", "Shake Shack"),
    ("https://upload.wikimedia.org/wikipedia/en/thumb/6/6d/Shakey%27s_US_logo.svg/330px-Shakey%27s_US_logo.svg.png", "Shakeys Pizza"),
    ("https://upload.wikimedia.org/wikipedia/commons/thumb/f/ff/SONIC_New_Logo_2020.svg/1199px-SONIC_New_Logo_2020.svg.png", "Sonic"),
    ("https://images.getbento.com/accounts/236fb3743b9522eafb90c6d2d20b8115/media/accounts/media/V1G5LTCrSDqyixbNUDpr_logo-tasty-burger.png", "Tasty burger"),
    ("https://cdn.worldvectorlogo.com/logos/uno-pizzeria.svg", "Uno Pizzeria"),
    ("https://upload.wikimedia.org/wikipedia/commons/thumb/2/28/Wahlburgers_Logo.svg/768px-Wahlburgers_Logo.svg.png", "Wahlburgers"),
    ("https://logowik.com/content/uploads/images/866_wendys.jpg", "Wendys"),
    ("https://logowik.com/content/uploads/images/whataburger8433.jpg", "Whataburger"),
    ("https://encrypted-tbn2.gstatic.com/images?q=tbn:ANd9GcTsUZPMFzUmACfaeSrMSZsgGn5gnbXNyYRPxIrS1PfCrFWI9QJhrlkmdG9IhFbi", "White Castle")]

    output = []
    for i in images_labels:
        if(i[1] in output_list):
          output.append(i)

    return output


def plot_heatmap():
  f = open("task1_output_n_c1_sup4.txt", "r")
  data = f.read()
  data = data.split("Frequent Itemsets:")

  clean_string = re.sub(r'\n', '', data[1])
  pattern = r"\([^\)]+\)"

  extracted_tuples = re.findall(pattern, clean_string.replace("'", ""))

  frequent_itemsets = [tuple(s.strip('()').split(', ')) for s in extracted_tuples]
  

  # Initialize a defaultdict for item pair frequencies
  pair_counts = defaultdict(int)

  # Count frequencies of item pairs, excluding self-pairings
  for itemset in frequent_itemsets:
      unique_items = list(set(itemset))
      for i in range(len(unique_items)):
          for j in range(i + 1, len(unique_items)):
              pair = tuple(sorted([unique_items[i], unique_items[j]]))
              pair_counts[pair] += 1

  # Get a sorted list of unique items
  all_items = sorted(set(item for pair in pair_counts for item in pair))

  # Create a DataFrame for the heatmap
  heatmap_data = pd.DataFrame(0, index=all_items, columns=all_items)
  for (item1, item2), count in pair_counts.items():
      heatmap_data.at[item1, item2] = count
      heatmap_data.at[item2, item1] = count  # Ensure symmetry

  # Create the heatmap
  fig = px.imshow(heatmap_data, labels=dict(x="Item", y="Item", color="Frequency"),
                  x=all_items, y=all_items,
                  color_continuous_scale=px.colors.sequential.Blues)

  fig.update_layout(title='Heatmap of Item Co-occurrences (Excluding Self-Pairings)',
                    title_x=0.5)

  return fig

def plot_bar():
  f = open("task1_output_n_c1_sup4.txt", "r")
  data = f.read()
  data = data.split("Frequent Itemsets:")

  clean_string = re.sub(r'\n', '', data[1])
  pattern = r"\([^\)]+\)"

  extracted_tuples = re.findall(pattern, clean_string.replace("'", ""))

  frequent_itemsets = [tuple(s.strip('()').split(', ')) for s in extracted_tuples]

  # Flatten the list of itemsets and count occurrences of each item
  items = [item for itemset in frequent_itemsets for item in itemset]
  item_counts = Counter(items)

  # Convert the counter to a DataFrame
  df = pd.DataFrame(item_counts.items(), columns=['Item', 'Count'])

  # Sort the DataFrame in descending order by 'Count'
  df = df.sort_values(by='Count', ascending=False)

  # Create the bar chart
  fig = px.bar(df, x='Item', y='Count', title='Frequency of Items in Frequent Itemsets',
              labels={'Count': 'Frequency'}, color='Count', height=600)

  fig.update_layout(xaxis_title='Item', yaxis_title='Frequency', title_x=0.5)

  return fig

def plot_association():
  f = open("task1_output_n_c1_sup4.txt", "r")
  data = f.read()
  data = data.split("Frequent Itemsets:")

  clean_string = re.sub(r'\n', '', data[1])
  pattern = r"\([^\)]+\)"

  extracted_tuples = re.findall(pattern, clean_string.replace("'", ""))

  frequent_itemsets = [tuple(s.strip('()').split(', ')) for s in extracted_tuples]

  G = nx.DiGraph()

  for itemset in frequent_itemsets:
      if len(itemset) == 2:
          G.add_edge(itemset[0], itemset[1])
      else:
          for i in range(len(itemset)):
              for j in range(i + 1, len(itemset)):
                  G.add_edge(itemset[i], itemset[j])

  pos = nx.spring_layout(G)

  edge_x = []
  edge_y = []
  for edge in G.edges():
      x0, y0 = pos[edge[0]]
      x1, y1 = pos[edge[1]]
      edge_x.append(x0)
      edge_x.append(x1)
      edge_x.append(None)
      edge_y.append(y0)
      edge_y.append(y1)
      edge_y.append(None)

  edge_trace = go.Scatter(
      x=edge_x, y=edge_y,
      line=dict(width=0.5, color='#888'),
      hoverinfo='none',
      mode='lines')

  node_x = []
  node_y = []
  for node in G.nodes():
      x, y = pos[node]
      node_x.append(x)
      node_y.append(y)

  node_trace = go.Scatter(
      x=node_x, y=node_y,
      mode='markers+text',
      text=[node for node in G.nodes()],
      textposition="top center",
      hoverinfo='text',
      marker=dict(
          showscale=True,
          colorscale='YlGnBu',
          size=10,
          colorbar=dict(
              thickness=15,
              title='Node Connections',
              xanchor='left',
              titleside='right'
          ),
          line_width=2))

  fig = go.Figure(data=[edge_trace, node_trace],
              layout=go.Layout(
                  title='<br>Association Rule Graph',
                  titlefont_size=16,
                  showlegend=False,
                  hovermode='closest',
                  margin=dict(b=20,l=5,r=5,t=40),
                  annotations=[ dict(
                      text="Association Rule Network",
                      showarrow=False,
                      xref="paper", yref="paper",
                      x=0.005, y=-0.002 ) ],
                  xaxis=dict(showgrid=False, zeroline=False),
                  yaxis=dict(showgrid=False, zeroline=False))
                  )
  return fig

sorted_unique_restaurants = [
 'A&W Restaurant',
 'Arbys',
 'Burger King',
 'California Kitchen',
 'Carls Jr.',
 'Chick-Fil-A',
 'Chuck E. Cheese',
 'Culvers',
 'Dairy Queen',
 'Dominos Pizza',
 'Five Guys',
 'Good Flippin Burger',
 'Hardees',
 'In-N-Out',
 'Jack in the Box',
 'Jollibee',
 'KFC',
 'Little Caesars',
 'McDonalds',
 'Papa Johns',
 'Pizza Inn',
 'Pizza hut',
 'PizzaExpress',
 'Popeyes',
 'Round Table Pizza',
 'Sbarro',
 'Shake Shack',
 'Shakeys Pizza',
 'Sonic',
 'Tasty burger',
 'Uno Pizzeria',
 'Wahlburgers',
 'Wendys',
 'Whataburger',
 'White Castle']

recommender = gr.Interface(fn=recommend, inputs=gr.Dropdown(choices=sorted_unique_restaurants, value="Burger King"), outputs=gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", columns=[4], object_fit="contain", height="auto"), title='Restaurant Recommender System', description='Find your ideal restaurant')  
vizz_heatmap = gr.Interface(fn=plot_heatmap, inputs = None, outputs=gr.Plot(label="Heatmap"), description = "Restaurant Itemset Frequency")
vizz_bar = gr.Interface(fn=plot_bar, inputs = None, outputs=gr.Plot(label="Bar Graph"), description = "Restaurant Frequency")
vizz_graph = gr.Interface(fn=plot_association, inputs = None, outputs=gr.Plot(label="Association Graph"), description = "Association Graph of Restaurants")


tabbed = gr.TabbedInterface([recommender, vizz_heatmap, vizz_bar, vizz_graph], ["Restaurant Recommender", "Heatmap", "Restaurant Frequency", "Association Graph"])

tabbed.queue().launch(share = True, debug = True)