app.py

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

# ============================================================
# 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
# ============================================================

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)
    company_to_sellers = defaultdict(list)
    company_to_buyers = defaultdict(list)

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

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

    for c in set(company_to_sellers.keys()) | set(company_to_buyers.keys()):
        sellers = company_to_sellers[c]
        buyers = company_to_buyers[c]
        for s in sellers:
            for b in buyers:
                transfers[(s, b)] += 1

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


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 u in G.nodes and v in G.nodes:
            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 s in G.nodes and b in G.nodes:
                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

# ============================================================
# PLOTLY NETWORK DRAWING
# ============================================================


def graph_to_plotly(
        G,
        node_color_amc="#9EC5FF",
        node_color_company="#FFCF9E",
        node_shape_amc="circle",
        node_shape_company="circle",
        edge_color_buy="#2ca02c",
        edge_color_sell="#d62728",
        edge_color_transfer="#888888",
        edge_thickness_base=1.4,
        show_labels=True
):
    pos = nx.spring_layout(G, seed=42, k=1.2)

    node_x, node_y, node_text, node_color, node_size = [], [], [], [], []

    for n, d in G.nodes(data=True):
        x, y = pos[n]
        node_x.append(x)
        node_y.append(y)
        node_text.append(n)

        if d["type"] == "amc":
            node_color.append(node_color_amc)
            node_size.append(40)
        else:
            node_color.append(node_color_company)
            node_size.append(60)

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

    edge_traces = []

    for u, v, attrs in G.edges(data=True):
        acts = attrs.get("actions", [])
        weight = attrs.get("weight", 1)

        x0, y0 = pos[u]
        x1, y1 = pos[v]

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

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

    fig = go.Figure(data=edge_traces + [node_trace])
    fig.update_layout(
        showlegend=False,
        height=900,
        width=1400,
        margin=dict(l=5, r=5, t=40, b=20),
        xaxis=dict(visible=False),
        yaxis=dict(visible=False)
    )
    return fig

# ============================================================
# COMPANY & AMC INSPECTION
# ============================================================


def company_trade_summary(company_name):
    buyers = [a for a, comps in BUY_MAP.items() if company_name in comps]
    sellers = [a for a, comps in SELL_MAP.items() if company_name in comps]
    fresh = [a for a, comps in FRESH_BUY.items() if company_name in comps]
    exits = [a for a, comps in COMPLETE_EXIT.items() if company_name in comps]

    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_text=f"Trade summary for {company_name}",
        height=300
    )
    return fig, df


def amc_transfer_summary(amc_name):
    sold = SELL_MAP.get(amc_name, [])
    transfers = []

    for s in sold:
        buyers = [a for a, comps in BUY_MAP.items() if s in comps]
        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="lightslategray"
    ))
    fig.update_layout(
        title_text=f"Inferred transfers from {amc_name}",
        height=300
    )
    return fig, df

# ============================================================
# INITIAL GRAPH
# ============================================================

initial_graph = build_graph(include_transfers=True)
initial_fig = graph_to_plotly(initial_graph)

# ============================================================
# GRADIO UI — CLEAN, FULL-WIDTH LAYOUT
# ============================================================

with gr.Blocks() as demo:
    gr.Markdown("## Mutual Fund Churn Explorer — Full Network & Transfer Analysis")

    # === FULL-WIDTH NETWORK GRAPH AT THE TOP ===
    network_plot = gr.Plot(
        value=initial_fig,
        label="Network graph (drag to zoom)"
    )

    # === SETTINGS BELOW THE GRAPH ===
    with gr.Accordion("Network Customization", open=True):
        node_color_company = gr.ColorPicker("#FFCF9E", label="Company node color")
        node_color_amc = gr.ColorPicker("#9EC5FF", label="AMC node color")

        node_shape_company = gr.Dropdown(["circle", "square", "diamond"], value="circle",
                                         label="Company node shape")
        node_shape_amc = gr.Dropdown(["circle", "square", "diamond"], value="circle",
                                     label="AMC node shape")

        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, 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")

    gr.Markdown("### Inspect a Company (buyers / sellers)")
    select_company = gr.Dropdown(choices=COMPANIES, label="Select company")
    company_out_plot = gr.Plot(label="Company trade summary")
    company_out_table = gr.DataFrame(label="Company table")

    gr.Markdown("### Inspect an AMC (transfer analysis)")
    select_amc = gr.Dropdown(choices=AMCS, label="Select AMC")
    amc_out_plot = gr.Plot(label="AMC transfer summary")
    amc_out_table = gr.DataFrame(label="AMC transfer table")

    # === CALLBACKS ===

    def update_network(node_color_company_val, node_color_amc_val,
                       node_shape_company_val, node_shape_amc_val,
                       edge_color_buy_val, edge_color_sell_val, edge_color_transfer_val,
                       edge_thickness_val, include_transfers_val):

        G = build_graph(include_transfers=include_transfers_val)
        fig = graph_to_plotly(
            G,
            node_color_amc=node_color_amc_val,
            node_color_company=node_color_company_val,
            node_shape_amc=node_shape_amc_val,
            node_shape_company=node_shape_company_val,
            edge_color_buy=edge_color_buy_val,
            edge_color_sell=edge_color_sell_val,
            edge_color_transfer=edge_color_transfer_val,
            edge_thickness_base=edge_thickness_val,
        )
        return fig

    update_button.click(
        update_network,
        [
            node_color_company,
            node_color_amc,
            node_shape_company,
            node_shape_amc,
            edge_color_buy,
            edge_color_sell,
            edge_color_transfer,
            edge_thickness,
            include_transfers,
        ],
        [network_plot]
    )

    def handle_company(company):
        fig, df = company_trade_summary(company)
        return fig, df

    def handle_amc(amc):
        fig, df = amc_transfer_summary(amc)
        return fig, df

    select_company.change(handle_company, select_company, [company_out_plot, company_out_table])
    select_amc.change(handle_amc, select_amc, [amc_out_plot, amc_out_table])


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