# app.py
# Mutual Fund Churn – Static Weighted Arc Diagram
# Short labels inside nodes, full labels on click
# Fully corrected version (NO syntax errors)

import gradio as gr
import pandas as pd
import json
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)

# ---------------------------
# Short + Full Label Maps
# ---------------------------
SHORT_LABELS = {
    "SBI MF": "SBI", "ICICI Pru MF": "ICICI", "HDFC MF": "HDFC",
    "Nippon India MF": "NIPP", "Kotak MF": "KOT", "UTI MF": "UTI",
    "Axis MF": "AXIS", "Aditya Birla SL MF": "ABSL", "Mirae MF": "MIR",
    "DSP MF": "DSP",

    "HDFC Bank": "HDFCB", "ICICI Bank": "ICICB",
    "Bajaj Finance": "BajFin", "Bajaj Finserv": "BajFsv",
    "Adani Ports": "AdPorts", "Tata Motors": "TataM",
    "Shriram Finance": "ShrFin", "HAL": "HAL", "TCS": "TCS",
    "AU Small Finance Bank": "AUSFB", "Pearl Global": "PearlG",
    "Hindalco": "Hind", "Tata Elxsi": "Elxsi",
    "Cummins India": "Cumm", "Vedanta": "Ved"
}

FULL_LABEL = {k: k for k in SHORT_LABELS}

# ---------------------------
# Infer AMC→AMC transfers
# ---------------------------
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) | set(c2b):
        for s in c2s[c]:
            for b in c2b[c]:
                transfers[(s, b)] += 1

    return transfers

TRANSFER_COUNTS = infer_amc_transfers(BUY_MAP, SELL_MAP)

# ---------------------------
# Mixed ordering for layout
# ---------------------------
def build_mixed_ordering(amcs, companies):
    out = []
    N = max(len(amcs), len(companies))
    for i in range(N):
        if i < len(amcs): out.append(amcs[i])
        if i < len(companies): out.append(companies[i])
    return out

NODES = build_mixed_ordering(AMCS, COMPANIES)
NODE_TYPE = {n: ("amc" if n in AMCS else "company") for n in NODES}

# ---------------------------
# Build flows
# ---------------------------
def build_flows():
    buys, sells, transfers, loops = [], [], [], []

    for amc, comps in BUY_MAP.items():
        for c in comps:
            w = 3 if c in FRESH_BUY.get(amc, []) else 1
            buys.append((amc, c, w))

    for amc, comps in SELL_MAP.items():
        for c in comps:
            w = 3 if c in COMPLETE_EXIT.get(amc, []) else 1
            sells.append((c, amc, w))

    for (s, b), w in TRANSFER_COUNTS.items():
        transfers.append((s, b, w))

    # loops (AMC → Company → AMC)
    seen = set()
    for a, c, _ in buys:
        for c2, b, _ in sells:
            if c == c2:
                seen.add((a, c, b))
    loops = list(seen)

    return buys, sells, transfers, loops

BUYS, SELLS, TRANSFERS, LOOPS = build_flows()

# ---------------------------
# Inspect panels
# ---------------------------
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, df

    counts = df.groupby("Role").size().reset_index(name="Count")
    fig = {
        "data": [{"type": "bar", "x": counts["Role"].tolist(), "y": counts["Count"].tolist()}],
        "layout": {"title": f"Trades for {company}"}
    }
    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, df

    counts = df["buyer_amc"].value_counts().reset_index()
    counts.columns = ["Buyer AMC", "Count"]
    fig = {
        "data": [{"type": "bar", "x": counts["Buyer AMC"].tolist(), "y": counts["Count"].tolist()}],
        "layout": {"title": f"Inferred transfers from {amc}"}
    }
    return fig, df

# ---------------------------
# JavaScript Template (raw, safe)
# ---------------------------
JS_TEMPLATE = r"""
<div id="arc-container" style="width:100%; height:720px;"></div>
<div style="margin-top:8px;">
  <button id="arc-reset" style="padding:8px 12px; border-radius:6px;">Reset</button>
</div>

<script src="https://d3js.org/d3.v7.min.js"></script>

<script>
const NODES = __NODES__;
const NODE_TYPE = __NODE_TYPE__;
const BUYS = __BUYS__;
const SELLS = __SELLS__;
const TRANSFERS = __TRANSFERS__;
const LOOPS = __LOOPS__;

const SHORT_LABEL_JS = __SHORT_LABEL__;
const FULL_LABEL_JS = __FULL_LABEL__;

function draw() {
    const container = document.getElementById("arc-container");
    container.innerHTML = "";

    const w = Math.min(950, container.clientWidth || 850);
    const h = Math.max(500, Math.floor(w * 0.75));
    const svg = d3.select(container).append("svg")
        .attr("width","100%")
        .attr("height",h)
        .attr("viewBox",[-w/2,-h/2,w,h].join(" "));

    const radius = Math.min(w,h)*0.36;

    const n = NODES.length;
    function angle(i){ return (i/n)*2*Math.PI; }

    const pos = NODES.map((name,i)=>{
        const ang = angle(i) - Math.PI/2;
        return {
            name,
            angle: ang,
            x: Math.cos(ang)*radius,
            y: Math.sin(ang)*radius
        };
    });

    const index = {};
    NODES.forEach((n,i)=> index[n]=i );

    // Node groups
    const g = svg.append("g").selectAll("g")
        .data(pos).enter().append("g")
        .attr("transform",d=>`translate(${d.x},${d.y})`);

    g.append("circle")
        .attr("r",16)
        .style("fill",d=> NODE_TYPE[d.name]==="amc" ? "#2b6fa6" : "#f2c88d")
        .style("stroke","#222")
        .style("stroke-width",1)
        .style("cursor","pointer");

    g.append("text")
        .attr("dy","0.35em")
        .style("font-size","9px")
        .style("fill","#fff")
        .style("text-anchor","middle")
        .style("pointer-events","none")
        .text(d=> SHORT_LABEL_JS[d.name] || d.name);

    // Helper for arcs
    function arcPath(x0,y0,x1,y1,above=true){
        const mx=(x0+x1)/2, my=(y0+y1)/2;
        const len=Math.sqrt(mx*mx+my*my)||1;
        const ux=mx/len, uy=my/len;
        const offset=(above?-1:1)*Math.max(35,radius*0.9);
        const cx=mx+ux*offset, cy=my+uy*offset;
        return `M ${x0} ${y0} Q ${cx} ${cy} ${x1} ${y1}`;
    }

    const allW = [].concat(
        BUYS.map(d=>d[2]), SELLS.map(d=>d[2]), TRANSFERS.map(d=>d[2])
    );
    const sw = d3.scaleLinear().domain([1,Math.max(...allW,1)]).range([1.2,6]);

    const buyG = svg.append("g");
    BUYS.forEach(([a,c,w])=>{
        if(!(a in index) || !(c in index)) return;
        const s=pos[index[a]], t=pos[index[c]];
        buyG.append("path")
            .attr("d",arcPath(s.x,s.y,t.x,t.y,true))
            .attr("stroke","#2e8540")
            .attr("fill","none")
            .attr("stroke-width",sw(w))
            .attr("data-src",a)
            .attr("data-tgt",c)
            .attr("opacity",0.92);
    });

    const sellG = svg.append("g");
    SELLS.forEach(([c,a,w])=>{
        if(!(c in index) || !(a in index)) return;
        const s=pos[index[c]], t=pos[index[a]];
        sellG.append("path")
            .attr("d",arcPath(s.x,s.y,t.x,t.y,false))
            .attr("stroke","#c0392b")
            .attr("fill","none")
            .attr("stroke-dasharray","4,3")
            .attr("stroke-width",sw(w))
            .attr("data-src",c)
            .attr("data-tgt",a)
            .attr("opacity",0.86);
    });

    const trG = svg.append("g");
    TRANSFERS.forEach(([s,b,w])=>{
        if(!(s in index)||!(b in index))return;
        const sp=pos[index[s]], tp=pos[index[b]];
        const mx=(sp.x+tp.x)/2, my=(sp.y+tp.y)/2;
        const path=`M ${sp.x} ${sp.y} Q ${mx*0.3} ${my*0.3} ${tp.x} ${tp.y}`;
        trG.append("path")
            .attr("d",path)
            .attr("stroke","#7d7d7d")
            .attr("fill","none")
            .attr("stroke-width",sw(w))
            .attr("opacity",0.7)
            .attr("data-src",s)
            .attr("data-tgt",b);
    });

    const loopG = svg.append("g");
    LOOPS.forEach(([a,c,b])=>{
        if(!(a in index)||!(b in index))return;
        const s=pos[index[a]], t=pos[index[b]];
        const mx=(s.x+t.x)/2, my=(s.y+t.y)/2;
        const len=Math.sqrt((s.x-t.x)**2+(s.y-t.y)**2);
        const offset=Math.max(40,radius*0.28+len*0.12);
        const ux=mx/Math.sqrt(mx*mx+my*my), uy=my/Math.sqrt(mx*mx+my*my);
        const cx=mx+ux*offset, cy=my+uy*offset;
        const path=`M ${s.x} ${s.y} Q ${cx} ${cy} ${t.x} ${t.y}`;
        loopG.append("path")
            .attr("d",path)
            .attr("stroke","#227a6d")
            .attr("fill","none")
            .attr("stroke-width",2.8)
            .attr("opacity",0.95);
    });

    function highlight(node){
        g.selectAll("circle").style("opacity",d=> d.name===node ? 1:0.18);
        g.selectAll("text")
            .style("opacity",d=> d.name===node ? 1:0.28)
            .text(d=> d.name===node ? FULL_LABEL_JS[d.name] : SHORT_LABEL_JS[d.name]);

        function match(el){ return el.getAttribute("data-src")===node || el.getAttribute("data-tgt")===node }

        buyG.selectAll("path").style("opacity",function(){return match(this)?0.98:0.06});
        sellG.selectAll("path").style("opacity",function(){return match(this)?0.98:0.06});
        trG.selectAll("path").style("opacity",function(){return match(this)?0.98:0.06});
    }

    function reset(){
        g.selectAll("circle").style("opacity",1);
        g.selectAll("text").style("opacity",1).text(d=>SHORT_LABEL_JS[d.name]);
        buyG.selectAll("path").style("opacity",0.92);
        sellG.selectAll("path").style("opacity",0.86);
        trG.selectAll("path").style("opacity",0.7);
        loopG.selectAll("path").style("opacity",0.95);
    }

    g.selectAll("circle").on("click",function(e,d){ highlight(d.name); e.stopPropagation(); });
    g.selectAll("text").on("click",function(e,d){ highlight(d.name); e.stopPropagation(); });

    document.getElementById("arc-reset").onclick=reset;
    svg.on("click",reset);
}

draw();
window.addEventListener("resize",draw);
</script>
"""

# ---------------------------
# Build HTML with replacements
# ---------------------------
def make_arc_html():
    html = JS_TEMPLATE
    html = html.replace("__NODES__", json.dumps(NODES))
    html = html.replace("__NODE_TYPE__", json.dumps(NODE_TYPE))
    html = html.replace("__BUYS__", json.dumps(BUYS))
    html = html.replace("__SELLS__", json.dumps(SELLS))
    html = html.replace("__TRANSFERS__", json.dumps(TRANSFERS))
    html = html.replace("__LOOPS__", json.dumps(LOOPS))
    html = html.replace("__SHORT_LABEL__", json.dumps(SHORT_LABELS))
    html = html.replace("__FULL_LABEL__", json.dumps(FULL_LABEL))
    return html

initial_html = make_arc_html()

# ---------------------------
# Gradio UI
# ---------------------------
with gr.Blocks(title="MF Churn – Arc Diagram") as demo:
    gr.Markdown("## Mutual Fund Churn — Weighted Arc Diagram (short labels → full label on click)")
    gr.HTML(initial_html)

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

    select_amc = gr.Dropdown(AMCS, label="Select AMC")
    amc_plot = gr.Plot()
    amc_table = gr.DataFrame()

    select_company.change(company_trade_summary, select_company, [company_plot, company_table])
    select_amc.change(amc_transfer_summary, select_amc, [amc_plot, amc_table])

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