# app.py
# Mutual Fund Churn Explorer with Gel + Wave Liquid Motion (Option D)
# D3 + Plotly hybrid layout
# Designed for Hugging Face Spaces (Gradio)

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)

# ============================================================
# GRAPH BUILDING
# ============================================================

def infer_amc_transfers(buy_map, sell_map):
    transfers = defaultdict(int)
    comp_sellers = defaultdict(list)
    comp_buyers = defaultdict(list)

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

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

    for c in set(comp_sellers.keys()) | set(comp_buyers.keys()):
        for s in comp_sellers[c]:
            for b in comp_buyers[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")

    # BUY/SELL edges
    for amc, comps in BUY_MAP.items():
        for c in comps:
            G.add_edge(amc, c, weight=1, actions=["buy"])

    for amc, comps in SELL_MAP.items():
        for c in comps:
            if G.has_edge(amc, c):
                G[amc][c]["weight"] += 1
                G[amc][c]["actions"].append("sell")
            else:
                G.add_edge(amc, c, weight=1, actions=["sell"])

    # complete exits
    for amc, comps in COMPLETE_EXIT.items():
        for c in comps:
            if G.has_edge(amc, c):
                G[amc][c]["weight"] += 3
                G[amc][c]["actions"].append("complete_exit")
            else:
                G.add_edge(amc, c, weight=3, actions=["complete_exit"])

    # fresh buys
    for amc, comps in FRESH_BUY.items():
        for c in comps:
            if G.has_edge(amc, c):
                G[amc][c]["weight"] += 3
                G[amc][c]["actions"].append("fresh_buy")
            else:
                G.add_edge(amc, c, weight=3, actions=["fresh_buy"])

    # inferred transfers
    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 FIGURE (placeholders — positions will be set by D3)
# ============================================================

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=1.4):

    node_names = []
    node_x = []
    node_y = []
    node_colors = []
    node_sizes = []

    for n, d in G.nodes(data=True):
        node_names.append(n)
        node_x.append(0)
        node_y.append(0)
        if d["type"] == "amc":
            node_colors.append(node_color_amc)
            node_sizes.append(36)
        else:
            node_colors.append(node_color_company)
            node_sizes.append(56)

    edge_traces = []
    src_idx = []
    tgt_idx = []
    e_colors = []
    e_widths = []

    for u, v, attrs in G.edges(data=True):
        edge_traces.append(
            go.Scatter(
                x=[0,0], y=[0,0], mode="lines",
                line=dict(color="#aaa", width=1),
                hoverinfo="none"
            )
        )
        src_idx.append(node_names.index(u))
        tgt_idx.append(node_names.index(v))

        acts = attrs.get("actions", [])
        w = attrs.get("weight", 1)

        if "complete_exit" in acts:
            e_colors.append(edge_color_sell)
            e_widths.append(edge_thickness * 3)
        elif "fresh_buy" in acts:
            e_colors.append(edge_color_buy)
            e_widths.append(edge_thickness * 3)
        elif "transfer" in acts:
            e_colors.append(edge_color_transfer)
            e_widths.append(edge_thickness * (1 + np.log1p(w)))
        elif "sell" in acts:
            e_colors.append(edge_color_sell)
            e_widths.append(edge_thickness * (1 + np.log1p(w)))
        else:
            e_colors.append(edge_color_buy)
            e_widths.append(edge_thickness * (1 + np.log1p(w)))

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

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

    meta = {
        "node_names": node_names,
        "edge_source_index": src_idx,
        "edge_target_index": tgt_idx,
        "edge_colors": e_colors,
        "edge_widths": e_widths,
        "node_sizes": node_sizes
    }

    return fig, meta


# ============================================================
# D3 + GEL + WAVE Motion Renderer
# ============================================================

def make_network_html(fig, meta, div_id="network-plot-div"):

    fig_json = json.dumps(fig.to_plotly_json())
    meta_json = json.dumps(meta)

    html = f"""
<div id="{div_id}" style="width:100%; height:620px;"></div>

<div style="margin-top:6px;">
  <button id="{div_id}-reset" style="padding:8px 12px; border-radius:6px;">Reset</button>
  <button id="{div_id}-stop" style="padding:8px 12px; margin-left:8px; border-radius:6px;">Stop Layout</button>
</div>

<script src="https://d3js.org/d3.v7.min.js"></script>
<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;

// Build nodes for D3
const nodes = meta.node_names.map((name, i) => {{
    return {{
        id: i,
        name: name,
        r: meta.node_sizes[i] || 20,
        displayX: 0,
        displayY: 0,
        vx_smooth: 0,
        vy_smooth: 0
    }};
}});

// Build links
const links = meta.edge_source_index.map((s, i) => {{
    return {{
        source: s,
        target: meta.edge_target_index[i],
        color: meta.edge_colors[i],
        width: meta.edge_widths[i]
    }};
}});

// D3 simulation
const simulation = d3.forceSimulation(nodes)
    .force("link", d3.forceLink(links).id(d => d.id).distance(150).strength(0.35))
    .force("charge", d3.forceManyBody().strength(-50))
    .force("collision", d3.forceCollide().radius(d => d.r * 0.9))
    .force("center", d3.forceCenter(0,0))
    .velocityDecay(0.55);

let tickCount = 0;
const maxTicks = 400;

simulation.on("tick", () => {{
    tickCount++;

    nodes.forEach(n => {{
        const tx = n.x || 0;
        const ty = n.y || 0;

        // Gel viscosity smoothing
        n.vx_smooth = n.vx_smooth * 0.82 + (tx - n.displayX) * 0.06;
        n.vy_smooth = n.vy_smooth * 0.82 + (ty - n.displayY) * 0.06;

        // Heavy damping
        n.vx_smooth *= 0.90;
        n.vy_smooth *= 0.90;

        // Update display positions
        n.displayX += n.vx_smooth;
        n.displayY += n.vy_smooth;

        // Wave pulse (gentle breathing)
        const t = Date.now() * 0.001;
        n.displayX += Math.sin(t + n.id) * 0.12;
        n.displayY += Math.cos(t + n.id) * 0.12;
    }});

    // Node arrays
    const xs = nodes.map(n => n.displayX);
    const ys = nodes.map(n => n.displayY);

    // Update node trace
    Plotly.restyle(container, {{x: [xs], y: [ys]}}, [nodeTraceIndex]);

    // Update edges
    for (let e = 0; e < edgeCount; e++) {{
        const s = meta.edge_source_index[e];
        const t = meta.edge_target_index[e];

        Plotly.restyle(container,
            {{
                x: [[nodes[s].displayX, nodes[t].displayX]],
                y: [[nodes[s].displayY, nodes[t].displayY]],
                "line.color": [meta.edge_colors[e]],
                "line.width": [meta.edge_widths[e]]
            }},
            [e]
        );
    }}

    if (simulation.alpha() < 0.02 || tickCount > maxTicks) {{
        simulation.stop();
    }}
}});

// STOP button
document.getElementById("{div_id}-stop").addEventListener("click", () => {{
    simulation.stop();
}});

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

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

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

    const N = meta.node_names.length;
    const op = Array(N).fill(0);
    const colors = Array(N).fill("rgba(0,0,0,0)");

    for (let i = 0; i < N; i++) {{
        if (keep.has(i)) {{
            op[i] = 1;
            colors[i] = "black";
        }}
    }}

    Plotly.restyle(container, {{
        "marker.opacity": [op],
        "textfont.color": [colors]
    }}, [nodeTraceIndex]);

    for (let e = 0; e < edgeCount; e++) {{
        const s = meta.edge_source_index[e], t = meta.edge_target_index[e];
        const show = keep.has(s) && keep.has(t);

        Plotly.restyle(container, {{
            "line.color": [show ? meta.edge_colors[e] : "rgba(0,0,0,0)"],
            "line.width": [show ? meta.edge_widths[e] : 0.1]
        }}, [e]);
    }}
}}

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

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

    simulation.alpha(0.5);
    simulation.restart();
}}

document.getElementById("{div_id}-reset").addEventListener("click", resetView);

container.on("plotly_click", (e) => {{
    const p = e.points[0];
    if (p && p.curveNumber === nodeTraceIndex) {{
        const name = meta.node_names[p.pointNumber];
        focusNode(name);
    }}
}});
</script>
"""

    return html

# ============================================================
# COMPANY / AMC SUMMARY
# ============================================================

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"Trades for {company}", margin=dict(t=30,b=5))

    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}", margin=dict(t=30,b=5))

    return fig, df


# ============================================================
# FINAL NETWORK HTML BUILDER
# ============================================================

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=edge_thickness
    )
    return make_network_html(fig, meta)

initial_html = build_network_html()


# ============================================================
# UI LAYOUT
# ============================================================

responsive_css = """
.js-plotly-plot { height:620px !important; }
@media(max-width:780px){ .js-plotly-plot{ height:600px !important; } }
"""

with gr.Blocks(css=responsive_css, title="MF Churn Explorer — Liquid Motion") as demo:

    gr.Markdown("## Mutual Fund Churn Explorer — Liquid Gel + Wave Motion (L2 + Rhythm)")

    network_html = gr.HTML(value=initial_html)

    legend_html = gr.HTML("""
<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 — inferred)</div>
  <div><span style="display:inline-block;width:28px;border-bottom:5px solid #2ca02c;"></span> FRESH BUY</div>
  <div><span style="display:inline-block;width:28px;border-bottom:5px solid #d62728;"></span> COMPLETE EXIT</div>
</div>
""")

    with gr.Accordion("Customize Network", 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.0, 1.4, step=0.1, label="Edge thickness")
        include_transfers = gr.Checkbox(True, label="Show inferred AMC→AMC transfers")
        update_btn = gr.Button("Update Graph")

    gr.Markdown("### Company Summary")
    select_company = gr.Dropdown(choices=COMPANIES, label="Select company")
    company_plot = gr.Plot()
    company_table = gr.DataFrame()

    gr.Markdown("### AMC Summary (Inferred Transfers)")
    select_amc = gr.Dropdown(choices=AMCS, label="Select AMC")
    amc_plot = gr.Plot()
    amc_table = gr.DataFrame()

    # Callbacks
    def update_net(c1,c2,buy,sell,trans,thick,inc):
        return build_network_html(
            node_color_company=c1,
            node_color_amc=c2,
            edge_color_buy=buy,
            edge_color_sell=sell,
            edge_color_transfer=trans,
            edge_thickness=thick,
            include_transfers=inc
        )

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

    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

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


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