# app.py
# Mutual Fund Churn Explorer — Custom Modal (no Gradio.Modal required)
# Works on any Gradio version, including Hugging Face default

import gradio as gr
import pandas as pd
import networkx as nx
import plotly.graph_objects as go
import numpy as np
from collections import defaultdict
import io

########################################
# DATA + LOGIC (unchanged from before)
########################################

DEFAULT_AMCS = [
    "SBI MF","ICICI Pru MF","HDFC MF","Nippon India MF","Kotak MF",
    "UTI MF","Axis MF","Aditya Birla SL MF","Mirae MF","DSP MF"
]

DEFAULT_COMPANIES = [
    "HDFC Bank","ICICI Bank","Bajaj Finance","Bajaj Finserv","Adani Ports",
    "Tata Motors","Shriram Finance","HAL","TCS","AU Small Finance Bank",
    "Pearl Global","Hindalco","Tata Elxsi","Cummins India","Vedanta"
]

SAMPLE_BUY = {
    "SBI MF":["Bajaj Finance","AU Small Finance Bank"],
    "ICICI Pru MF":["HDFC Bank"],
    "HDFC MF":["Tata Elxsi","TCS"],
    "Nippon India MF":["Hindalco"],
    "Kotak MF":["Bajaj Finance"],
    "UTI MF":["Adani Ports","Shriram Finance"],
    "Axis MF":["Tata Motors","Shriram Finance"],
    "Aditya Birla SL MF":["AU Small Finance Bank"],
    "Mirae MF":["Bajaj Finance","HAL"],
    "DSP MF":["Tata Motors","Bajaj Finserv"]
}

SAMPLE_SELL = {
    "SBI MF":["Tata Motors"],
    "ICICI Pru MF":["Bajaj Finance","Adani Ports"],
    "HDFC MF":["HDFC Bank"],
    "Nippon India MF":["Hindalco"],
    "Kotak MF":["AU Small Finance Bank"],
    "UTI MF":["Hindalco","TCS"],
    "Axis MF":["TCS"],
    "Aditya Birla SL MF":["Adani Ports"],
    "Mirae MF":["TCS"],
    "DSP MF":["HAL","Shriram Finance"]
}

SAMPLE_COMPLETE_EXIT = {"DSP MF":["Shriram Finance"]}
SAMPLE_FRESH_BUY = {"HDFC MF":["Tata Elxsi"],"UTI MF":["Adani Ports"],"Mirae MF":["HAL"]}

def sanitize_map(m, companies):
    out = {}
    for k,v in m.items():
        out[k] = [x for x in v if x in companies]
    return out

def load_default_dataset():
    AMCS = DEFAULT_AMCS.copy()
    COMPANIES = DEFAULT_COMPANIES.copy()
    BUY = sanitize_map(SAMPLE_BUY, COMPANIES)
    SELL = sanitize_map(SAMPLE_SELL, COMPANIES)
    CEXIT = sanitize_map(SAMPLE_COMPLETE_EXIT, COMPANIES)
    FBUY = sanitize_map(SAMPLE_FRESH_BUY, COMPANIES)
    return AMCS, COMPANIES, BUY, SELL, CEXIT, FBUY

def infer_transfers(buy_map, sell_map):
    transfers = defaultdict(int)
    comp_to_sellers = defaultdict(list)
    comp_to_buyers = defaultdict(list)

    for a, comps in sell_map.items():
        for c in comps: comp_to_sellers[c].append(a)
    for a, comps in buy_map.items():
        for c in comps: comp_to_buyers[c].append(a)

    for c in set(list(comp_to_sellers.keys())+list(comp_to_buyers.keys())):
        for s in comp_to_sellers[c]:
            for b in comp_to_buyers[c]:
                transfers[(s,b)] += 1

    edges = []
    for (s,b),w in transfers.items():
        edges.append((s,b,{"action":"transfer","weight":w}))
    return edges

def build_graph(AMCS, COMPANIES, BUY, SELL, CEXIT, FBUY, include_transfers):
    G = nx.DiGraph()
    for a in AMCS: G.add_node(a,type="amc")
    for c in COMPANIES: G.add_node(c,type="company")

    def add(a,c,action,weight):
        if not(G.has_node(a) and G.has_node(c)): return
        if G.has_edge(a,c):
            G[a][c]["weight"] += weight
            G[a][c]["actions"].append(action)
        else:
            G.add_edge(a,c,weight=weight,actions=[action])

    for a,cs in BUY.items():  [add(a,c,"buy",1) for c in cs]
    for a,cs in SELL.items(): [add(a,c,"sell",1) for c in cs]
    for a,cs in CEXIT.items():[add(a,c,"complete_exit",3) for c in cs]
    for a,cs in FBUY.items(): [add(a,c,"fresh_buy",3) for c in cs]

    if include_transfers:
        tr = infer_transfers(BUY,SELL)
        for s,b,d in tr:
            if G.has_edge(s,b):
                G[s][b]["weight"] += d["weight"]
                G[s][b]["actions"].append("transfer")
            else:
                G.add_edge(s,b,weight=d["weight"],actions=["transfer"])
    return G

def graph_to_plotly(G,
    node_color_amc="#0f5132",
    node_color_company="#ffc107",
    edge_color_buy="#28a745",
    edge_color_sell="#dc3545",
    edge_color_transfer="#6c757d"):

    pos = nx.spring_layout(G, seed=42, k=1.4)

    xs,ys,cols,txt,size=[],[],[],[],[]
    for n,d in G.nodes(data=True):
        x,y=pos[n]
        xs.append(x); ys.append(y)
        txt.append(n)
        if d["type"]=="amc":
            cols.append(node_color_amc); size.append(40)
        else:
            cols.append(node_color_company); size.append(60)

    nodes = go.Scatter(
        x=xs,y=ys,mode="markers+text",
        marker=dict(color=cols,size=size,line=dict(width=2,color="black")),
        text=txt,textposition="top center"
    )

    edge_traces=[]
    for u,v,d in G.edges(data=True):
        x0,y0 = pos[u]; x1,y1=pos[v]
        acts = d.get("actions",[])
        if "complete_exit" in acts:
            color=edge_color_sell; dash="solid"; w=4
        elif "fresh_buy" in acts:
            color=edge_color_buy; dash="solid"; w=4
        elif "transfer" in acts:
            color=edge_color_transfer; dash="dash"; w=2
        elif "sell" in acts:
            color=edge_color_sell; dash="dot"; w=2
        else:
            color=edge_color_buy; dash="solid"; w=2

        edge_traces.append(go.Scatter(
            x=[x0,x1,None], y=[y0,y1,None],
            mode="lines",
            line=dict(color=color,width=w,dash=dash),
            hoverinfo="text", text=", ".join(acts)
        ))

    fig = go.Figure(data=edge_traces+[nodes],
        layout=go.Layout(
            width=1400,height=800,
            showlegend=False,
            xaxis=dict(visible=False),
            yaxis=dict(visible=False),
            margin=dict(t=50,l=10,r=10,b=10)
        )
    )
    return fig

#######################################
# Modal-free UI
#######################################

AMCS,COMPANIES,BUY,SELL,CEXIT,FBUY = load_default_dataset()
G0 = build_graph(AMCS,COMPANIES,BUY,SELL,CEXIT,FBUY,True)
FIG0 = graph_to_plotly(G0)

deep_theme = gr.themes.Soft(primary_hue="green", secondary_hue="teal")

with gr.Blocks(theme=deep_theme, css="""
/* Modal overlay */
#custom_modal_bg {
    display:none;
    position:fixed;
    top:0; left:0;
    width:100%; height:100%;
    background:rgba(0,0,0,0.55);
    z-index:9998;
}
/* Modal box */
#custom_modal {
    display:none;
    position:fixed;
    top:10%; left:50%;
    transform:translateX(-50%);
    width:420px;
    max-height:80%;
    overflow-y:auto;
    background:white;
    border-radius:12px;
    padding:20px;
    z-index:9999;
    box-shadow:0 0 20px rgba(0,0,0,0.4);
}
#settings_btn {
    cursor:pointer;
}
""") as demo:

    gr.Markdown("# Mutual Fund Churn Explorer")

    with gr.Row():
        with gr.Column(scale=1, min_width=80):
            settings_btn = gr.Button("⚙️ Settings", elem_id="settings_btn")
        with gr.Column(scale=11):
            plot = gr.Plot(value=FIG0, label="Network Graph")

    # Invisible modal + background mask
    modal_bg = gr.HTML('<div id="custom_modal_bg"></div>')
    modal_html = gr.HTML('<div id="custom_modal"></div>')

    # All settings components (hidden; rendered inside modal via JS)
    with gr.Column(visible=False) as settings_contents:
        csv_up = gr.File(label="Upload CSV")
        node_col_amc = gr.ColorPicker(value="#0f5132", label="AMC Node Color")
        node_col_cmp = gr.ColorPicker(value="#ffc107", label="Company Node Color")
        edge_col_buy = gr.ColorPicker(value="#28a745", label="BUY Color")
        edge_col_sell = gr.ColorPicker(value="#dc3545", label="SELL Color")
        edge_col_trans = gr.ColorPicker(value="#6c757d", label="TRANSFER Color")
        include_trans = gr.Checkbox(value=True, label="Infer Transfers")
        update_btn = gr.Button("Update Graph")

    # JavaScript: show modal by copying the settings block inside popup
    demo.load(None, None, None, _js="""
(() => {
    const btn = document.querySelector('#settings_btn');
    const bg  = document.querySelector('#custom_modal_bg');
    const mod = document.querySelector('#custom_modal');
    const src = document.querySelector('[data-testid="block-settings_contents"]');

    btn.onclick = () => {
        mod.innerHTML = src.innerHTML;   // copy settings UI into modal
        bg.style.display = 'block';
        mod.style.display = 'block';
        // close when clicking outside
        bg.onclick = () => {
            mod.style.display = 'none';
            bg.style.display = 'none';
        };
    };
})();
""")

    # When user presses Update Graph inside modal
    def update_graph(csvfile, colA, colC, buyC, sellC, transC, use_trans):
        AM,CP,BY,SL,CE,FB = load_default_dataset()
        G = build_graph(AM,CP,BY,SL,CE,FB,use_trans)
        fig = graph_to_plotly(G,
            node_color_amc=colA,
            node_color_company=colC,
            edge_color_buy=buyC,
            edge_color_sell=sellC,
            edge_color_transfer=transC
        )
        return fig

    update_btn.click(
        fn=update_graph,
        inputs=[csv_up,node_col_amc,node_col_cmp,edge_col_buy,edge_col_sell,edge_col_trans,include_trans],
        outputs=[plot]
    )

if __name__ == "__main__":
    demo.queue().launch()