// Copyright (c) 2026 Ranjith N
// SPDX-License-Identifier: MIT OR Apache-2.0

use csv::Reader;
use serde::Deserialize;
use std::collections::HashMap;
use std::error::Error;
use std::fs::File;
use svg::node::element::{Circle, Line, Rectangle, Text};
use svg::node::element::Group;
use svg::Document;

#[derive(Debug)]
struct Args {
    input_file: String,
    meta_filt: String,
    output_file: String,
}


#[derive(Debug, Deserialize)]
struct Record {
    meta: String,
    #[serde(rename = "Name1")]
    name1: String,
    #[serde(rename = "Name2")]
    name2: String,
    vardirect: String,
    slope: f64,
    weighted_c_score: f64,
}

pub fn run(args: Vec<String>) -> Result<(), Box<dyn Error>> {

    // Show usage if no args
    if args.len() < 7 {
        eprintln!("Usage: {} -i <input.csv> -m <meta_filt> -o <output.svg>", args[0]);
        std::process::exit(1);
    }

    // Manual argument parsing
    let mut input_file = String::new();
    let mut meta_filt = String::new();
    let mut output_file = String::new();

    let mut i = 1;
    while i < args.len() {
        match args[i].as_str() {
            "-i" => {
                input_file = args.get(i + 1).cloned().unwrap_or_default();
                i += 1;
            }
            "-m" => {
                meta_filt = args.get(i + 1).cloned().unwrap_or_default();
                i += 1;
            }
            "-o" => {
                output_file = args.get(i + 1).cloned().unwrap_or_default();
                i += 1;
            }
            _ => {}
        }
        i += 1;
    }

    // Check required flags
    if input_file.is_empty() || meta_filt.is_empty() || output_file.is_empty() {
        eprintln!("❌ Missing required arguments.");
        eprintln!("Usage: {} -i <input.csv> -m <meta_filt> -o <output.svg>", args[0]);
        std::process::exit(1);
    }

    // Pack into struct
    let args = Args {
        input_file,
        meta_filt,
        output_file,
    };

    let file = File::open(&args.input_file)?;
    let mut rdr = Reader::from_reader(file);

    let mut edge_data = Vec::new();

    for result in rdr.deserialize() {
        let record: Record = result?;

        if record.meta != args.meta_filt {
            continue;
        }
        if record.weighted_c_score == 0.0 || record.vardirect == "??" {
            continue;
        }

        edge_data.push(record);
    }

    let mut degree_count: HashMap<String, usize> = HashMap::new();
    for edge in &edge_data {
        *degree_count.entry(edge.name1.clone()).or_insert(0) += 1;
        *degree_count.entry(edge.name2.clone()).or_insert(0) += 1;
    }

    let mut top_nodes: Vec<_> = degree_count.iter().collect();
    top_nodes.sort_by(|a, b| b.1.cmp(a.1).then_with(|| a.0.cmp(b.0)));
    let top_nodes: Vec<String> = top_nodes.into_iter().take(30).map(|(k, _)| k.clone()).collect();

    let node_idx: HashMap<String, usize> = top_nodes
        .iter()
        .enumerate()
        .map(|(i, name)| (name.clone(), i))
        .collect();

    let filtered_edges: Vec<_> = edge_data
        .into_iter()
        .filter(|e| node_idx.contains_key(&e.name1) && node_idx.contains_key(&e.name2))
        .collect();

    plot_bubble_chart(&filtered_edges, &top_nodes, &node_idx, &args)?;

    Ok(())
}

fn plot_bubble_chart(
    edges: &[Record],
    nodes: &[String],
    node_idx: &HashMap<String, usize>,
    args: &Args,
) -> Result<(), Box<dyn Error>> {
    let plot_size = 1000;

    // Margins for axis ticks and labels
    let margin_left = 40;
    let margin_bottom = 40;
    let margin_top = 30;
//    let margin_right = 20;

    // Calculate dynamic legend space based on longest node name
    let max_name_len = nodes.iter().map(|s| s.len()).max().unwrap_or(10);
    let avg_char_width = 8; // average pixel width per character (adjust if needed)
    let legend_space = (max_name_len as i32 * avg_char_width + 150).max(300);

    let canvas_width = plot_size + margin_left + legend_space;
    let canvas_height = plot_size + margin_top + margin_bottom;

    let center_x = margin_left + plot_size / 2;
    let center_y = margin_top + plot_size / 2;

    let n = nodes.len() as i32;
    let padding = 2;
    let axis_limit = n + padding;

    let scale = (plot_size / 2) as f64 / (axis_limit as f64); // scale per axis unit

    // Min/max for bubble size normalization
    let roeffect_min = edges.iter().map(|e| e.slope).fold(f64::INFINITY, f64::min);
    let roeffect_max = edges.iter().map(|e| e.slope).fold(f64::NEG_INFINITY, f64::max);

    let norm = |val: f64| -> i32 {
        let min_size = 1.0;
        let max_size = 15.0;
        let normalized = (val - roeffect_min) / (roeffect_max - roeffect_min + 1e-9);
        (min_size + (max_size - min_size) * normalized).round() as i32
    };

    let green = "#34A853";   // example hex color for green
    let orange = "#FBBC04";  // example hex color for orange
    let blue = "#4285F4";    // example hex color for blue
    let red = "#EA4335";     // example hex color for red
    
    let color_map = HashMap::from([
        ("+/+".to_string(), green),
        ("+/-".to_string(), orange),
        ("-/+".to_string(), blue),
        ("-/-".to_string(), red),
    ]);

    let mut document = Document::new()
        .set("width", canvas_width)
        .set("height", canvas_height)
        .set("viewBox", (0, 0, canvas_width, canvas_height));

    // Draw main plot rectangle (shifted by margins)
    let border = Rectangle::new()
        .set("x", margin_left)
        .set("y", margin_top)
        .set("width", plot_size)
        .set("height", plot_size)
        .set("fill", "white")
        .set("stroke", "black")
        .set("stroke-width", 1);
    document = document.add(border);

    // Add plot title (meta filter value) at the top center
    let title_text = Text::new()
        .set("x", margin_left + plot_size / 2)
        .set("y", margin_top / 2 + 10)  // place it in the top margin area, and 10 is to avoid
                                        // clipping
        .set("text-anchor", "middle")
        .set("font-family", "sans-serif")
        .set("font-size", 20)
        .set("font-weight", "bold")
        .add(svg::node::Text::new(&args.meta_filt));
    
    document = document.add(title_text);


    // Draw internal grid lines
    for i in -axis_limit..=axis_limit {
        let pos = (i as f64 * scale).round() as i32;

        // Vertical grid lines
        document = document.add(Line::new()
            .set("x1", center_x + pos)
            .set("y1", margin_top)
            .set("x2", center_x + pos)
            .set("y2", margin_top + plot_size)
            .set("stroke", "#ccc")
            .set("stroke-width", 1));

        // Horizontal grid lines
        document = document.add(Line::new()
            .set("x1", margin_left)
            .set("y1", center_y - pos)
            .set("x2", margin_left + plot_size)
            .set("y2", center_y - pos)
            .set("stroke", "#ccc")
            .set("stroke-width", 1));
    }

    // Draw central axes lines (inside plot, shifted by margins)
    document = document.add(Line::new()
        .set("x1", margin_left)
        .set("y1", center_y)
        .set("x2", margin_left + plot_size)
        .set("y2", center_y)
        .set("stroke", "black")
        .set("stroke-width", 1));

    document = document.add(Line::new()
        .set("x1", center_x)
        .set("y1", margin_top)
        .set("x2", center_x)
        .set("y2", margin_top + plot_size)
        .set("stroke", "black")
        .set("stroke-width", 1));

    // Draw edges and bubbles
    for edge in edges {
        let i = *node_idx.get(&edge.name1).unwrap() as i32;
        let j = *node_idx.get(&edge.name2).unwrap() as i32;

        let (x, y) = match edge.vardirect.as_str() {
            "+/+" => (i + 1, j + 1),
            "-/-" => (-i - 1, -j - 1),
            "+/-" => (-j - 1, i + 1),
            "-/+" => (j + 1, -i - 1),
            _ => continue,
        };

        let svg_x = center_x as i32 + (x as f64 * scale).round() as i32;
        let svg_y = center_y as i32 - (y as f64 * scale).round() as i32;

        let color = color_map.get(&edge.vardirect).unwrap_or(&"grey");
        let size = norm(edge.slope);

        let line_x = Line::new()
            .set("x1", center_x)
            .set("y1", svg_y)
            .set("x2", svg_x)
            .set("y2", svg_y)
            .set("stroke", *color)
            .set("stroke-width", 1);
        let line_y = Line::new()
            .set("x1", svg_x)
            .set("y1", center_y)
            .set("x2", svg_x)
            .set("y2", svg_y)
            .set("stroke", *color)
            .set("stroke-width", 1);

        let circle = Circle::new()
            .set("cx", svg_x)
            .set("cy", svg_y)
            .set("r", size)
            .set("fill", *color); // no stroke

        document = document.add(line_x).add(line_y).add(circle);
    }

    // Draw Y axis ticks and labels on LEFT (in margin)
    let tick_length = 6;
    let font_size = 12;
    for i in -axis_limit..=axis_limit {
        let y = center_y - (i as f64 * scale).round() as i32;

        // tick line at left edge of plot
        document = document.add(Line::new()
            .set("x1", margin_left)
            .set("y1", y)
            .set("x2", margin_left - tick_length)
            .set("y2", y)
            .set("stroke", "black")
            .set("stroke-width", 1));

        // label to the left of tick inside margin
        document = document.add(Text::new()
            .set("x", margin_left - tick_length - 4) // 4px padding from tick
            .set("y", y + font_size / 3)
            .set("text-anchor", "end")
            .set("font-family", "sans-serif")
            .set("font-size", font_size)
            .add(svg::node::Text::new(format!("{}", i))));
    }

    // Draw X axis ticks and labels on BOTTOM (in margin)
    for i in -axis_limit..=axis_limit {
        let x = center_x + (i as f64 * scale).round() as i32;

        // tick line at bottom edge of plot
        document = document.add(Line::new()
            .set("x1", x)
            .set("y1", margin_top + plot_size)
            .set("x2", x)
            .set("y2", margin_top + plot_size + tick_length)
            .set("stroke", "black")
            .set("stroke-width", 1));

        // label below the tick
        document = document.add(Text::new()
            .set("x", x)
            .set("y", margin_top + plot_size + tick_length + font_size)
            .set("text-anchor", "middle")
            .set("font-family", "sans-serif")
            .set("font-size", font_size)
            .add(svg::node::Text::new(format!("{}", i))));
    }

    // Draw legend group (NO background box)
    let legend_x_offset = margin_left + plot_size + 10;
    let mut legend_group = Group::new()
        .set("transform", format!("translate({}, 0)", legend_x_offset));

    // Draw node labels with numbering
    for (i, name) in nodes.iter().enumerate() {
        let text = Text::new()
            .set("x", 10)
            .set("y", 20 + i as i32 * 20)
            .set("font-family", "sans-serif")
            .set("font-size", 15)
            .add(svg::node::Text::new(format!("{:2}: {}", i + 1, name)));
        legend_group = legend_group.add(text);
    }

//    let bubble_start_y = 20 + nodes.len() as i32 * 20 + 40;
    let num_bubbles = 5;
    let bubble_x_center = 90;
    let bubble_spacing = 30;

    // Bubble size legend below nodes
    let bubble_legend_height = (num_bubbles as i32 * bubble_spacing) + 40; // includes padding
    let quadrant_legend_height = 150; // estimate or calculate dynamically if needed
    
    // Total vertical offset from bottom
//    let total_legend_height = bubble_legend_height + quadrant_legend_height + 20; // 20 padding
    
    let legend_bottom_y = canvas_height as i32 - 20; // 20px padding from bottom
    
    // bubble legend starts just above the quadrant legend
    let bubble_start_y = legend_bottom_y - quadrant_legend_height - bubble_legend_height;


    for i in 0..num_bubbles {
        let slope = roeffect_min + (roeffect_max - roeffect_min) * (i as f64) / ((num_bubbles - 1) as f64);
        let size = norm(slope);
        let y = bubble_start_y + i as i32 * bubble_spacing;

        let circle = Circle::new()
            .set("cx", bubble_x_center)
            .set("cy", y)
            .set("r", size)
            .set("fill", "grey");

        let text = Text::new()
            .set("x", bubble_x_center + size + 10)
            .set("y", y + (size / 2))
            .set("font-family", "sans-serif")
            .set("font-size", 12)
            .add(svg::node::Text::new(format!("{:.2}", slope)));

        legend_group = legend_group.add(circle).add(text);
    }

    // ───────────── Quadrant Legend (Below Bubble Legend) ─────────────
    
    use svg::node::element::{Path, Marker, Definitions};
    use svg::node::element::path::Data;
    
    // Arrowhead marker definition
    let arrow_marker = Marker::new()
        .set("id", "arrowhead")
        .set("markerWidth", 6)
        .set("markerHeight", 6)
        .set("refX", 3)
        .set("refY", 3)
        .set("orient", "auto")
        .set("markerUnits", "strokeWidth")
        .add(Path::new()
            .set("d", "M0,0 L5,3 L0,6 L2.5,3 z")
            .set("fill", "black")); // Use black for visibility or adjust as needed
    
    // Add arrow marker to defs (attach only once to the document)
    let defs = Definitions::new().add(arrow_marker);
    document = document.add(defs);
    
    // Positioning for the quadrant legend
    let quad_legend_start_y = legend_bottom_y - quadrant_legend_height;
//    let quad_legend_start_y = bubble_start_y + num_bubbles as i32 * bubble_spacing + 50;
    let quad_scale = 110.0; // scale unit to fit small quadrant diagram
    let quad_center_x = 90.0;
    let quad_center_y = quad_legend_start_y as f64 + 60.0;
    
    // Axes
    let axis_x = Line::new()
        .set("x1", quad_center_x - 0.6 * quad_scale)
        .set("y1", quad_center_y)
        .set("x2", quad_center_x + 0.6 * quad_scale)
        .set("y2", quad_center_y)
        .set("stroke", "black")
        .set("stroke-width", 1);
    
    let axis_y = Line::new()
        .set("x1", quad_center_x)
        .set("y1", quad_center_y - 0.6 * quad_scale)
        .set("x2", quad_center_x)
        .set("y2", quad_center_y + 0.6 * quad_scale)
        .set("stroke", "black")
        .set("stroke-width", 1);
    
    // Arcs with arrowheads
    let arcs = vec![
        ("#34A853", 0.60, 0.10, 0.10, 0.60),    // Q1 green
        ("#FBBC04", -0.10, 0.60, -0.60, 0.10),  // Q2 orange
        ("#EA4335", -0.60, -0.10, -0.10, -0.60),// Q3 red
        ("#4285F4", 0.10, -0.60, 0.60, -0.10),  // Q4 blue
    ];
    
    for (color, x1, y1, x2, y2) in arcs {
        let data = Data::new()
            .move_to((
                quad_center_x + x1 * quad_scale,
                quad_center_y - y1 * quad_scale,
            ))
            .elliptical_arc_to((
                0.55 * quad_scale,
                0.55 * quad_scale,
                0.0,
                0,
                0,
                quad_center_x + x2 * quad_scale,
                quad_center_y - y2 * quad_scale,
            ));
    
        let arc = Path::new()
            .set("d", data)
            .set("stroke", color)
            .set("stroke-width", 4)
            .set("fill", "none")
            .set("marker-end", "url(#arrowhead)");
    
        legend_group = legend_group.add(arc);
    }
    
    // Labels
    let labels = vec![
        ("+x", 0.65, 0.0, "start"),
        ("-x", -0.65, 0.0, "end"),
        ("+y", 0.0, 0.65, "middle"),
        ("-y", 0.0, -0.65, "middle"),
    ];
    
    for (text, dx, dy, anchor) in labels {
        let label = Text::new()
            .set("x", quad_center_x + dx * quad_scale)
            .set("y", quad_center_y - dy * quad_scale)
            .set("fill", "black")
            .set("font-weight", "bold")
            .set("font-size", 14)
            .set("dominant-baseline", "middle")
            .set("text-anchor", anchor)
            .add(svg::node::Text::new(text));
    
        legend_group = legend_group.add(label);
    }
    
    // Quadrant labels
    let quadrant_labels = vec![
        ("Q1 (+/+)", 0.25, 0.25, "#34A853"),
        ("Q2 (+/-)", -0.25, 0.25, "#FBBC04"),
        ("Q3 (-/-)", -0.25, -0.25, "#EA4335"),
        ("Q4 (-/+)", 0.25, -0.25, "#4285F4"),
    ];
    
    for (text, dx, dy, color) in quadrant_labels {
        let label = Text::new()
            .set("x", quad_center_x + dx * quad_scale)
            .set("y", quad_center_y - dy * quad_scale)
            .set("fill", color)
            .set("font-weight", "bold")
            .set("font-size", 14)
            .set("text-anchor", "middle")
            .set("dominant-baseline", "middle")
            .add(svg::node::Text::new(text));
    
        legend_group = legend_group.add(label);
    }
    
    // Add axes last to draw on top
    legend_group = legend_group.add(axis_x).add(axis_y);


    // Add legend group to document
    document = document.add(legend_group);

    // Save SVG
    svg::save(&args.output_file, &document)?;

    println!("SVG saved to {}", args.output_file);
    Ok(())
}

