# app.py
# Interactive MF churn explorer — with client-side clickable Plotly
# NOW WITH LEGEND UNDER CHART (only addition requested)

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

# ---------------------------
# Data
# ---------------------------
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"
]

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"
]

BUY_MAP = {
    "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"]
}

SELL_MAP = {
    "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"]
}

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

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

BUY_MAP = sanitize_map(BUY_MAP)
SELL_MAP = sanitize_map(SELL_MAP)
COMPLETE_EXIT = sanitize_map(COMPLETE_EXIT)
FRESH_BUY = sanitize_map(FRESH_BUY)

# ---------------------------
# Build graph edges
# ---------------------------
company_edges = []
for amc, comps in BUY_MAP.items():
    for c in comps:
        company_edges.append((amc, c, {"action": "buy", "weight": 1}))
for amc, comps in SELL_MAP.items():
    for c in comps:
        company_edges.append((amc, c, {"action": "sell", "weight": 1}))
for amc, comps in COMPLETE_EXIT.items():
    for c in comps:
        company_edges.append((amc, c, {"action": "complete_exit", "weight": 3}))
for amc, comps in FRESH_BUY.items():
    for c in comps:
        company_edges.append((amc, c, {"action": "fresh_buy", "weight": 3}))

def infer_amc_transfers(buy_map, sell_map):
    transfers = defaultdict(int)
    c2s = defaultdict(list)
    c2b = defaultdict(list)

    for amc, comps in sell_map.items():
        for c in comps:
            c2s[c].append(amc)

    for amc, comps in buy_map.items():
        for c in comps:
            c2b[c].append(amc)

    for c in set(c2s.keys()) | set(c2b.keys()):
        for s in c2s[c]:
            for b in c2b[c]:
                transfers[(s,b)] += 1

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

transfer_edges = infer_amc_transfers(BUY_MAP, SELL_MAP)

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

    for u,v,attr in company_edges:
        if G.has_edge(u,v):
            G[u][v]["weight"] += attr["weight"]
            G[u][v]["actions"].append(attr["action"])
        else:
            G.add_edge(u,v,weight=attr["weight"], actions=[attr["action"]])

    if include_transfers:
        for s,b,attr in transfer_edges:
            if G.has_edge(s,b):
                G[s][b]["weight"] += attr["weight"]
                G[s][b]["actions"].append("transfer")
            else:
                G.add_edge(s,b,weight=attr["weight"], actions=["transfer"])

    return G

# ---------------------------
# Build Plotly figure for embedding
# ---------------------------
def build_plotly_figure(G,
                        node_color_amc="#9EC5FF",
                        node_color_company="#FFCF9E",
                        edge_color_buy="#2ca02c",
                        edge_color_sell="#d62728",
                        edge_color_transfer="#888888",
                        edge_thickness_base=1.4):

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

    node_names = []
    node_x = []
    node_y = []
    node_color = []
    node_size = []

    for n, d in G.nodes(data=True):
        node_names.append(n)
        x, y = pos[n]
        node_x.append(x); node_y.append(y)
        if d["type"] == "amc":
            node_color.append(node_color_amc); node_size.append(36)
        else:
            node_color.append(node_color_company); node_size.append(56)

    edge_traces = []
    edge_source_index = []
    edge_target_index = []
    edge_colors = []
    edge_widths = []

    for u,v,attrs in G.edges(data=True):
        x0,y0 = pos[u]; x1,y1 = pos[v]
        acts = attrs["actions"]
        weight = attrs["weight"]

        if "complete_exit" in acts:
            color = edge_color_sell; width = edge_thickness_base*3; dash="solid"
        elif "fresh_buy" in acts:
            color = edge_color_buy; width = edge_thickness_base*3; dash="solid"
        elif "transfer" in acts:
            color = edge_color_transfer; width=edge_thickness_base*(1+np.log1p(weight)); dash="dash"
        elif "sell" in acts:
            color = edge_color_sell; width=edge_thickness_base*(1+np.log1p(weight)); dash="dot"
        else:
            color = edge_color_buy; width=edge_thickness_base*(1+np.log1p(weight)); dash="solid"

        edge_traces.append(go.Scatter(
            x=[x0,x1], y=[y0,y1],
            mode="lines",
            line=dict(color=color, width=width, dash=dash),
            hoverinfo="none",
            opacity=1.0
        ))
        edge_source_index.append(node_names.index(u))
        edge_target_index.append(node_names.index(v))
        edge_colors.append(color)
        edge_widths.append(width)

    node_trace = go.Scatter(
        x=node_x, y=node_y,
        mode="markers+text",
        marker=dict(color=node_color, size=node_size, line=dict(width=2,color="#222")),
        text=node_names,
        textposition="top center",
        hoverinfo="text"
    )

    fig = go.Figure(data=edge_traces+[node_trace])
    fig.update_layout(
        showlegend=False,
        autosize=True,
        margin=dict(l=8,r=8,t=36,b=8),
        xaxis=dict(visible=False),
        yaxis=dict(visible=False)
    )

    meta = {
        "node_names": node_names,
        "edge_source_index": edge_source_index,
        "edge_target_index": edge_target_index,
        "edge_colors": edge_colors,
        "edge_widths": edge_widths
    }
    return fig, meta

# ---------------------------
# Create HTML with JS click-to-focus behavior
# ---------------------------
def make_network_html(fig, meta, div_id="network-plot-div"):
    fig_json = json.dumps(fig.to_plotly_json())
    meta_json = json.dumps(meta)

    return f"""
<div id="{div_id}" style="width:100%;height:520px;"></div>
<div style="margin-top:6px;margin-bottom:8px;">
  <button id="{div_id}-reset" style="padding:8px 12px;border-radius:6px;">Reset view</button>
</div>

<script src="https://cdn.plot.ly/plotly-latest.min.js"></script>

<script>
const fig = {fig_json};
const meta = {meta_json};

const container = document.getElementById("{div_id}");

Plotly.newPlot(container, fig.data, fig.layout, {{responsive:true}});

const nodeTraceIndex = fig.data.length - 1;
const edgeCount = fig.data.length - 1;

const nameToIndex = {{}};
meta.node_names.forEach((n,i)=>nameToIndex[n]=i);

function focusNode(nodeName){{
    const idx = nameToIndex[nodeName];
    const keep = new Set([idx]);

    for(let e=0;e<meta.edge_source_index.length;e++){{
        const s=meta.edge_source_index[e];
        const t=meta.edge_target_index[e];
        if(s===idx) keep.add(t);
        if(t===idx) keep.add(s);
    }}

    
    // Update nodes (hide others + hide their labels)
    const N = meta.node_names.length;
    const nodeOp = Array(N).fill(0.0);
    const textColors = Array(N).fill("rgba(0,0,0,0)");
    
    for (let i = 0; i < N; i++) {
        if (keep.has(i)) {
            nodeOp[i] = 1.0;
            textColors[i] = "black";   // visible label
        }
    }
    
    Plotly.restyle(container, {
        "marker.opacity": [nodeOp],
        "textfont.color": [textColors]
    }, [nodeTraceIndex]);


    // edges
    for(let e=0;e<edgeCount;e++){{
        const s=meta.edge_source_index[e];
        const t=meta.edge_target_index[e];
        const show = keep.has(s)&&keep.has(t);
        const color = show?meta.edge_colors[e]:"rgba(0,0,0,0)";
        const width = show?meta.edge_widths[e]:0.1;
        Plotly.restyle(container,{{
            "line.color":[color],
            "line.width":[width]
        }},[e]);
    }}
}}

function resetView(){{
    const N=meta.node_names.length;
    const op=Array(N).fill(1.0);

    const N = meta.node_names.length;
    const nodeOp = Array(N).fill(1.0);
    const textColors = Array(N).fill("black");
    
    Plotly.restyle(container, {
        "marker.opacity":[nodeOp],
        "textfont.color":[textColors]
    }, [nodeTraceIndex]);


    for(let e=0;e<edgeCount;e++){{
        Plotly.restyle(container,{{
            "line.color":[meta.edge_colors[e]],
            "line.width":[meta.edge_widths[e]]
        }},[e]);
    }}

    Plotly.relayout(container, {{
        xaxis: {{autorange:true}},
        yaxis: {{autorange:true}}
    }});
}}

container.on('plotly_click', function(evt){{
    const p = evt.points[0];
    if(p.curveNumber===nodeTraceIndex){{
        const idx = p.pointNumber;
        const name = meta.node_names[idx];
        focusNode(name);
    }}
}});

document.getElementById("{div_id}-reset").onclick = resetView;
</script>
"""

# ---------------------------------------------
# Build HTML network block
# ---------------------------------------------
def build_network_html(node_color_company="#FFCF9E",
                       node_color_amc="#9EC5FF",
                       edge_color_buy="#2ca02c",
                       edge_color_sell="#d62728",
                       edge_color_transfer="#888888",
                       edge_thickness=1.4,
                       include_transfers=True):

    G = build_graph(include_transfers=include_transfers)
    fig, meta = build_plotly_figure(
        G,
        node_color_amc=node_color_amc,
        node_color_company=node_color_company,
        edge_color_buy=edge_color_buy,
        edge_color_sell=edge_color_sell,
        edge_color_transfer=edge_color_transfer,
        edge_thickness_base=edge_thickness
    )
    return make_network_html(fig, meta)

# Initial HTML
initial_html = build_network_html()

# ---------------------------
# Company & AMC summaries
# ---------------------------
def company_trade_summary(company):
    buyers = [a for a,cs in BUY_MAP.items() if company in cs]
    sellers = [a for a,cs in SELL_MAP.items() if company in cs]
    fresh = [a for a,cs in FRESH_BUY.items() if company in cs]
    exits = [a for a,cs in COMPLETE_EXIT.items() if company in cs]

    df = pd.DataFrame({
        "Role": (["Buyer"]*len(buyers)) + (["Seller"]*len(sellers)) +
                (["Fresh buy"]*len(fresh)) + (["Complete exit"]*len(exits)),
        "AMC": buyers + sellers + fresh + exits
    })

    if df.empty:
        return None, pd.DataFrame(columns=["Role","AMC"])

    counts = df.groupby("Role").size().reset_index(name="Count")
    fig = go.Figure(go.Bar(x=counts["Role"], y=counts["Count"],
                           marker_color=["green","red","orange","black"][:len(counts)]))
    fig.update_layout(title=f"Trade summary for {company}", autosize=True)
    return fig, df

def amc_transfer_summary(amc):
    sold = SELL_MAP.get(amc, [])
    transfers = []
    for s in sold:
        buyers = [a for a,cs in BUY_MAP.items() if s in cs]
        for b in buyers:
            transfers.append({"security": s, "buyer_amc": b})
    df = pd.DataFrame(transfers)
    if df.empty:
        return None, pd.DataFrame(columns=["security","buyer_amc"])

    counts = df["buyer_amc"].value_counts().reset_index()
    counts.columns = ["Buyer AMC","Count"]
    fig = go.Figure(go.Bar(x=counts["Buyer AMC"], y=counts["Count"], marker_color="gray"))
    fig.update_layout(title=f"Inferred transfers from {amc}", autosize=True)
    return fig, df

# ---------------------------
# Mobile-friendly CSS
# ---------------------------
responsive_css = """
.gradio-container { padding:0 !important; margin:0 !important; }
.plotly-graph-div, .js-plotly-plot { width:100% !important; max-width:100% !important; }
.js-plotly-plot { height:460px !important; }
@media(max-width:780px){ .js-plotly-plot{ height:420px !important; } }
body, html { overflow-x:hidden !important; }
"""

# ---------------------------
# UI BLOCKS WITH LEGEND ADDED
# ---------------------------
with gr.Blocks(css=responsive_css, title="MF Churn Explorer") as demo:

    gr.Markdown("## Mutual Fund Churn Explorer — Interactive Graph")

    # Chart (interactive HTML)
    network_html = gr.HTML(value=initial_html)

    # ⭐ LEGEND (ONLY addition)
    legend_html = gr.HTML(value="""
<div style='
    font-family: sans-serif;
    font-size: 14px;
    margin-top: 10px;
    line-height: 1.6;
'>
    <b>Legend</b><br>

    <div>
        <span style="display:inline-block;width:28px;
            border-bottom:3px solid #2ca02c;"></span>
        BUY (green solid)
    </div>

    <div>
        <span style="display:inline-block;width:28px;
            border-bottom:3px dotted #d62728;"></span>
        SELL (red dotted)
    </div>

    <div>
        <span style="display:inline-block;width:28px;
            border-bottom:3px dashed #888;"></span>
        TRANSFER (grey dashed)
    </div>

    <div>
        <span style="display:inline-block;width:28px;
            border-bottom:5px solid #2ca02c;"></span>
        FRESH BUY (thick green)
    </div>

    <div>
        <span style="display:inline-block;width:28px;
            border-bottom:5px solid #d62728;"></span>
        COMPLETE EXIT (thick red)
    </div>
</div>
""")

    # Controls
    with gr.Accordion("Network Customization — expand to edit", open=False):
        node_color_company = gr.ColorPicker("#FFCF9E", label="Company node color")
        node_color_amc = gr.ColorPicker("#9EC5FF", label="AMC node color")
        edge_color_buy = gr.ColorPicker("#2ca02c", label="BUY edge color")
        edge_color_sell = gr.ColorPicker("#d62728", label="SELL edge color")
        edge_color_transfer = gr.ColorPicker("#888888", label="Transfer edge color")
        edge_thickness = gr.Slider(0.5, 6, value=1.4, step=0.1, label="Edge thickness base")
        include_transfers = gr.Checkbox(value=True, label="Show AMC→AMC inferred transfers")
        update_button = gr.Button("Update Network Graph")

    # Company inspect
    gr.Markdown("### Inspect Company (buyers / sellers)")
    select_company = gr.Dropdown(COMPANIES, label="Select company")
    company_plot = gr.Plot()
    company_table = gr.DataFrame()

    # AMC inspect
    gr.Markdown("### Inspect AMC (inferred transfers)")
    select_amc = gr.Dropdown(AMCS, label="Select AMC")
    amc_plot = gr.Plot()
    amc_table = gr.DataFrame()

    # Callbacks
    def update_network(node_color_company_val,
                       node_color_amc_val,
                       edge_color_buy_val,
                       edge_color_sell_val,
                       edge_color_transfer_val,
                       edge_thickness_val,
                       include_transfers_val):
        return build_network_html(
            node_color_company=node_color_company_val,
            node_color_amc=node_color_amc_val,
            edge_color_buy=edge_color_buy_val,
            edge_color_sell=edge_color_sell_val,
            edge_color_transfer=edge_color_transfer_val,
            edge_thickness=edge_thickness_val,
            include_transfers=include_transfers_val
        )

    def on_company(c):
        fig, df = company_trade_summary(c)
        return fig, df

    def on_amc(a):
        fig, df = amc_transfer_summary(a)
        return fig, df

    update_button.click(
        update_network,
        inputs=[node_color_company, node_color_amc,
                edge_color_buy, edge_color_sell, edge_color_transfer,
                edge_thickness, include_transfers],
        outputs=[network_html]
    )

    select_company.change(on_company, [select_company], [company_plot, company_table])
    select_amc.change(on_amc, [select_amc], [amc_plot, amc_table])

# Run app
if __name__ == "__main__":
    demo.launch()