use std::fs::File;
use std::io::{BufRead, BufReader};

use cch_rs::cch::*;
use petgraph::graph::DiGraph;

pub struct DimacsGraph {
    pub first_out: Vec<u32>,
    pub head: Vec<u32>,
    pub weights: Vec<f32>,
    pub xs: Vec<f32>,
    pub ys: Vec<f32>,
}

impl CchGraph for DimacsGraph {
    fn num_nodes(&self) -> usize { self.first_out.len() - 1 }
    fn num_edges(&self) -> usize { self.head.len() }
    fn first_out(&self) -> &[u32] { &self.first_out }
    fn head(&self) -> &[u32] { &self.head }
    fn x(&self, node_id: u32) -> f32 { self.xs[node_id as usize] }
    fn y(&self, node_id: u32) -> f32 { self.ys[node_id as usize] }
}

pub struct DimacsLoader;

impl DimacsLoader {
    pub fn load_dimacs(gr_path: &str, co_path: &str) -> Result<DimacsGraph, Box<dyn std::error::Error>> {
        let mut raw_xs = Vec::new();
        let mut raw_ys = Vec::new();
        let mut node_ids = Vec::new();

        let co_file = BufReader::new(File::open(co_path)?);
        for line in co_file.lines().flatten() {
            let p: Vec<&str> = line.split_whitespace().collect();
            if p.get(0) == Some(&"v") {
                node_ids.push(p[1].parse::<u32>()?);
                raw_xs.push(p[2].parse::<f32>()?);
                raw_ys.push(p[3].parse::<f32>()?);
            }
        }

        let mut temp_edges = Vec::new();
        let gr_file = BufReader::new(File::open(gr_path)?);
        for line in gr_file.lines().flatten() {
            let p: Vec<&str> = line.split_whitespace().collect();
            if p.get(0) == Some(&"a") {
                let u = p[1].parse::<u32>()?;
                let v = p[2].parse::<u32>()?;
                let w = p[3].parse::<f32>()?;
                temp_edges.push((u, v, w));
            }
        }

        let mut sorted_unique_ids = node_ids.clone();
        sorted_unique_ids.sort_unstable();

        let get_internal_id = |original_id: u32| -> usize { sorted_unique_ids.binary_search(&original_id).expect("Node ID not found") };

        let num_nodes = sorted_unique_ids.len();
        let mut xs = vec![0.0; num_nodes];
        let mut ys = vec![0.0; num_nodes];

        for (i, &oid) in node_ids.iter().enumerate() {
            let internal_id = get_internal_id(oid);
            xs[internal_id] = raw_xs[i];
            ys[internal_id] = raw_ys[i];
        }

        let mut mapped_edges: Vec<(u32, u32, f32)> = temp_edges.into_iter().map(|(u, v, w)| (get_internal_id(u) as u32, get_internal_id(v) as u32, w)).collect();

        mapped_edges.sort_unstable_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1)));

        let mut first_out = Vec::with_capacity(num_nodes + 1);
        let mut head = Vec::with_capacity(mapped_edges.len());
        let mut weights = Vec::with_capacity(mapped_edges.len());

        let mut edge_iter = mapped_edges.into_iter().peekable();
        for u in 0..num_nodes as u32 {
            first_out.push(head.len() as u32);
            while let Some(edge) = edge_iter.peek() {
                if edge.0 == u {
                    let (_, v, w) = edge_iter.next().unwrap();
                    if !head.is_empty() && *head.last().unwrap() == v && u == u {
                        let last_idx = weights.len() - 1;
                        if w < weights[last_idx] {
                            weights[last_idx] = w;
                        }
                    } else {
                        head.push(v);
                        weights.push(w);
                    }
                } else {
                    break;
                }
            }
        }
        first_out.push(head.len() as u32);

        Ok(DimacsGraph { first_out, head, weights, xs, ys })
    }

    pub fn to_petgraph(graph: &DimacsGraph) -> DiGraph<(), f32> {
        let mut g = DiGraph::new();
        let nodes: Vec<_> = (0..graph.num_nodes()).map(|_| g.add_node(())).collect();
        for u in 0..graph.num_nodes() as u32 {
            let start = graph.first_out[u as usize] as usize;
            let end = graph.first_out[u as usize + 1] as usize;
            for i in start..end {
                g.add_edge(nodes[u as usize], nodes[graph.head[i] as usize], graph.weights[i]);
            }
        }
        g
    }
}
