# app.py
# Static weighted semi-layer arc diagram (L1 labels outside)
# JS block injected safely (no Python f-strings inside JS)

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)

# ---------------------------
# 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.keys()) | set(c2b.keys()):
        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 to reduce crossings
# ---------------------------
def build_mixed_ordering(amcs, companies):
    mixed = []
    n = max(len(amcs), len(companies))
    for i in range(n):
        if i < len(amcs):
            mixed.append(amcs[i])
        if i < len(companies):
            mixed.append(companies[i])
    return mixed

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 = []
    for amc, comps in BUY_MAP.items():
        for c in comps:
            w = 3 if (amc in FRESH_BUY and c in FRESH_BUY.get(amc, [])) else 1
            buys.append((amc, c, w))
    sells = []
    for amc, comps in SELL_MAP.items():
        for c in comps:
            w = 3 if (amc in COMPLETE_EXIT and c in COMPLETE_EXIT.get(amc, [])) else 1
            sells.append((c, amc, w))
    transfers = []
    for (s,b), w in TRANSFER_COUNTS.items():
        transfers.append((s, b, w))
    loops = []
    # loops: a -> c -> b
    for a,c,w1 in buys:
        for c2,b,w2 in sells:
            if c == c2:
                loops.append((a, c, b))
    # dedupe
    loops = list({(a,c,b) for (a,c,b) in loops})
    return buys, sells, transfers, loops

BUYS, SELLS, TRANSFERS, LOOPS = build_flows()

# ---------------------------
# Inspector 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 = {
        "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, pd.DataFrame([], columns=["security", "buyer_amc"])
    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

# ---------------------------
# HTML template (JS inserted safely via replace)
# ---------------------------
JS_TEMPLATE = """
<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>

<div style="margin-top:10px; font-family:sans-serif; font-size:13px;">
  <b>Legend</b><br/>
  <span style="display:inline-block;width:12px;height:8px;background:#2e8540;margin-right:6px;"></span> BUY (green solid)<br/>
  <span style="display:inline-block;width:12px;height:8px;background:#c0392b;margin-right:6px;border-bottom:3px dotted #c0392b;"></span> SELL (red dotted)<br/>
  <span style="display:inline-block;width:12px;height:8px;background:#7d7d7d;margin-right:6px;"></span> TRANSFER (grey, inferred)<br/>
  <span style="display:inline-block;width:12px;height:8px;background:#227a6d;margin-right:6px;"></span> LOOP (external arc)<br/>
  <div style="margin-top:6px;color:#666;font-size:12px;">Note: Transfers are inferred by matching sells and buys across AMCs. Thickness shows relative weight.</div>
</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__ ;

function draw() {
    const container = document.getElementById("arc-container");
    container.innerHTML = "";
    const w = Math.min(920, container.clientWidth || 820);
    const h = Math.max(420, 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;

    // node positions around circle
    const n = NODES.length;
    function angleFor(i) { return (i / n) * 2 * Math.PI; }
    const nodePos = NODES.map((name,i) => {
        const ang = angleFor(i) - Math.PI/2; // start at top
        return { name: name, angle: ang, x: Math.cos(ang)*radius, y: Math.sin(ang)*radius };
    });

    const nameToIndex = {};
    NODES.forEach((nm,i)=> nameToIndex[nm]=i);

    // small node circles and labels outside
    const group = svg.append("g").selectAll("g").data(nodePos).enter().append("g")
        .attr("transform", d => `translate(${d.x},${d.y})`);

    group.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");

    group.append("text")
        .attr("x", d => Math.cos(d.angle) * (radius + 26))
        .attr("y", d => Math.sin(d.angle) * (radius + 26))
        .attr("dy", "0.35em")
        .style("font-family", "sans-serif")
        .style("font-size", Math.max(10, Math.min(14, radius*0.04)))
        .style("text-anchor", d => {
            const deg = (d.angle * 180 / Math.PI);
            return (deg > -90 && deg < 90) ? "start" : "end";
        })
        .attr("transform", d => {
            const deg = (d.angle * 180 / Math.PI);
            const flip = (deg > 90 || deg < -90) ? 180 : 0;
            return `rotate(${deg}) translate(${radius + 26}) rotate(${flip})`;
        })
        .style("cursor","pointer")
        .text(d => d.name);

    // bezier helper
    function bezierPath(x0,y0,x1,y1,above=true) {
        const mx = (x0 + x1)/2;
        const my = (y0 + y1)/2;
        const dx = mx;
        const dy = my;
        const len = Math.sqrt(dx*dx + dy*dy) || 1;
        const ux = dx/len, uy = dy/len;
        const offset = (above ? -1 : 1) * Math.max(30, radius*0.9);
        const cx = mx + ux * offset;
        const cy = my + uy * offset;
        return `M ${x0} ${y0} Q ${cx} ${cy} ${x1} ${y1}`;
    }

    // stroke width scale
    const allW = [].concat(BUYS.map(d=>d[2]), SELLS.map(d=>d[2]), TRANSFERS.map(d=>d[2]));
    const wmin = Math.min(...(allW.length?allW:[1]));
    const wmax = Math.max(...(allW.length?allW:[1]));
    const stroke = d3.scaleLinear().domain([wmin, Math.max(wmax,1)]).range([1.0, 6.0]);

    // buys (top)
    const buyGroup = svg.append("g").attr("class","buys");
    BUYS.forEach(b => {
        const a = b[0], c = b[1], wt = b[2];
        if (!(a in nameToIndex) || !(c in nameToIndex)) return;
        const s = nodePos[nameToIndex[a]];
        const t = nodePos[nameToIndex[c]];
        const path = bezierPath(s.x,s.y,t.x,t.y,true);
        buyGroup.append("path")
            .attr("d", path)
            .attr("fill","none")
            .attr("stroke","#2e8540")
            .attr("stroke-width", stroke(wt))
            .attr("stroke-linecap","round")
            .attr("opacity", 0.92)
            .attr("data-src", a)
            .attr("data-tgt", c)
            .on("mouseover", function() { d3.select(this).attr("opacity",1); })
            .on("mouseout", function() { d3.select(this).attr("opacity",0.92); });
    });

    // sells (bottom)
    const sellGroup = svg.append("g").attr("class","sells");
    SELLS.forEach(s => {
        const c = s[0], a = s[1], wt = s[2];
        if (!(c in nameToIndex) || !(a in nameToIndex)) return;
        const sp = nodePos[nameToIndex[c]];
        const tp = nodePos[nameToIndex[a]];
        const path = bezierPath(sp.x,sp.y,tp.x,tp.y,false);
        sellGroup.append("path")
            .attr("d", path)
            .attr("fill","none")
            .attr("stroke","#c0392b")
            .attr("stroke-width", stroke(wt))
            .attr("stroke-linecap","round")
            .attr("stroke-dasharray","4,3")
            .attr("opacity",0.86)
            .attr("data-src", c)
            .attr("data-tgt", a)
            .on("mouseover", function() { d3.select(this).attr("opacity",1); })
            .on("mouseout", function() { d3.select(this).attr("opacity",0.86); });
    });

    // transfers (grey chords)
    const transferGroup = svg.append("g").attr("class","transfers");
    TRANSFERS.forEach(tr => {
        const sname = tr[0], tname = tr[1], wt = tr[2];
        if (!(sname in nameToIndex) || !(tname in nameToIndex)) return;
        const sp = nodePos[nameToIndex[sname]];
        const tp = nodePos[nameToIndex[tname]];
        const mx = (sp.x + tp.x)/2;
        const my = (sp.y + tp.y)/2;
        const cx = mx * 0.3, cy = my * 0.3;
        const path = `M ${sp.x} ${sp.y} Q ${cx} ${cy} ${tp.x} ${tp.y}`;
        transferGroup.append("path")
            .attr("d", path)
            .attr("fill","none")
            .attr("stroke","#7d7d7d")
            .attr("stroke-width", stroke(wt))
            .attr("opacity",0.7)
            .attr("data-src", sname)
            .attr("data-tgt", tname)
            .on("mouseover", function() { d3.select(this).attr("opacity",1); })
            .on("mouseout", function() { d3.select(this).attr("opacity",0.7); });
    });

    // loops (external arcs)
    const loopGroup = svg.append("g").attr("class","loops");
    LOOPS.forEach(lp => {
        const a = lp[0], c = lp[1], b = lp[2];
        if (!(a in nameToIndex) || !(b in nameToIndex)) return;
        const sa = nodePos[nameToIndex[a]];
        const sb = nodePos[nameToIndex[b]];
        const mx = (sa.x + sb.x)/2;
        const my = (sa.y + sb.y)/2;
        const len = Math.sqrt((sa.x - sb.x)*(sa.x - sb.x) + (sa.y - sb.y)*(sa.y - sb.y));
        const outward = Math.max(40, radius*0.28 + len * 0.12);
        const ndx = mx, ndy = my;
        const nlen = Math.sqrt(ndx*ndx + ndy*ndy) || 1;
        const ux = ndx/nlen, uy = ndy/nlen;
        const cx = mx + ux * outward;
        const cy = my + uy * outward;
        const path = `M ${sa.x} ${sa.y} Q ${cx} ${cy} ${sb.x} ${sb.y}`;
        loopGroup.append("path")
            .attr("d", path)
            .attr("fill", "none")
            .attr("stroke", "#227a6d")
            .attr("stroke-width", 2.8)
            .attr("opacity",0.95)
            .on("mouseover", function() { d3.select(this).attr("opacity",1); })
            .on("mouseout", function() { d3.select(this).attr("opacity",0.95); });
    });

    // interactivity: focus on node
    function setOpacityFor(nodeName) {
        group.selectAll("circle").style("opacity", d => (d.name === nodeName ? 1.0 : 0.18));
        group.selectAll("text").style("opacity", d => (d.name === nodeName ? 1.0 : 0.28));
        buyGroup.selectAll("path").style("opacity", function() {
            const src = this.getAttribute("data-src");
            const tgt = this.getAttribute("data-tgt");
            return (src === nodeName || tgt === nodeName) ? 0.98 : 0.06;
        });
        sellGroup.selectAll("path").style("opacity", function() {
            const src = this.getAttribute("data-src");
            const tgt = this.getAttribute("data-tgt");
            return (src === nodeName || tgt === nodeName) ? 0.98 : 0.06;
        });
        transferGroup.selectAll("path").style("opacity", function() {
            const src = this.getAttribute("data-src");
            const tgt = this.getAttribute("data-tgt");
            return (src === nodeName || tgt === nodeName) ? 0.98 : 0.06;
        });
        loopGroup.selectAll("path").style("opacity", 0.95);
    }

    function resetOpacity() {
        group.selectAll("circle").style("opacity", 1.0);
        group.selectAll("text").style("opacity", 1.0);
        buyGroup.selectAll("path").style("opacity", 0.92);
        sellGroup.selectAll("path").style("opacity", 0.86);
        transferGroup.selectAll("path").style("opacity", 0.7);
        loopGroup.selectAll("path").style("opacity", 0.95);
    }

    group.selectAll("circle").style("cursor","pointer").on("click", function(e,d) {
        setOpacityFor(d.name);
        if (e && e.stopPropagation) e.stopPropagation();
    });
    group.selectAll("text").style("cursor","pointer").on("click", function(e,d) {
        setOpacityFor(d.name);
        if (e && e.stopPropagation) e.stopPropagation();
    });

    document.getElementById("arc-reset").onclick = resetOpacity;
    svg.on("click", function(event) {
        if (event.target.tagName === "svg") resetOpacity();
    });
}

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

def make_arc_html(nodes, node_type, buys, sells, transfers, loops):
    # prepare JSON strings
    nodes_json = json.dumps(nodes)
    node_type_json = json.dumps(node_type)
    buys_json = json.dumps(buys)
    sells_json = json.dumps(sells)
    transfers_json = json.dumps(transfers)
    loops_json = json.dumps(loops)
    html = JS_TEMPLATE.replace("__NODES__", nodes_json) \
                      .replace("__NODE_TYPE__", node_type_json) \
                      .replace("__BUYS__", buys_json) \
                      .replace("__SELLS__", sells_json) \
                      .replace("__TRANSFERS__", transfers_json) \
                      .replace("__LOOPS__", loops_json)
    return html

initial_html = make_arc_html(NODES, NODE_TYPE, BUYS, SELLS, TRANSFERS, LOOPS)

# ---------------------------
# Gradio UI
# ---------------------------
responsive_css = """
#arc-container { padding:0; margin:0; }
svg { font-family: sans-serif; }
"""

with gr.Blocks(css=responsive_css, title="MF Churn — Semi-layer Arc Diagram (L1 labels)") as demo:
    gr.Markdown("## Mutual Fund Churn — Weighted Arcs (BUY top / SELL bottom) — labels outside (L1)")
    gr.HTML(initial_html)

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

    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()