use crate::{PyGenotype, PyOp, bindings::gp::PyGraph};
use pyo3::{Bound, IntoPyObjectExt, PyAny, PyResult, Python, pyclass, pymethods};
use radiate::{Codec, GraphCodec, NodeType, Op};
use std::collections::HashMap;

const INPUT_NODE_TYPE: &str = "input";
const OUTPUT_NODE_TYPE: &str = "output";
const VERTEX_NODE_TYPE: &str = "vertex";
const EDGE_NODE_TYPE: &str = "edge";

#[pyclass(from_py_object)]
#[derive(Clone)]
pub struct PyGraphCodec {
    pub codec: GraphCodec<Op<f32>>,
}

#[pymethods]
impl PyGraphCodec {
    pub fn encode_py(&self) -> PyResult<PyGenotype> {
        Ok(PyGenotype::from(self.codec.encode()))
    }

    pub fn decode_py<'py>(
        &self,
        py: Python<'py>,
        genotype: &PyGenotype,
    ) -> PyResult<Bound<'py, PyAny>> {
        let genotype = genotype.clone().into();
        let obj_value = self.codec.decode(&genotype);

        PyGraph {
            inner: obj_value,
            eval_cache: None,
        }
        .into_bound_py_any(py)
    }

    #[new]
    #[pyo3(signature = (graph_type=None, input_size=1, output_size=1, ops=None, max_nodes=None))]
    pub fn new<'py>(
        graph_type: Option<&'py str>,
        input_size: usize,
        output_size: usize,
        ops: Option<HashMap<String, Vec<PyOp>>>,
        max_nodes: Option<usize>,
    ) -> PyResult<Self> {
        let mut values = Vec::new();
        if let Some(ops) = &ops {
            for (key, value) in ops.iter() {
                if key == INPUT_NODE_TYPE {
                    values.push((
                        NodeType::Input,
                        value.iter().map(|op| op.0.clone()).collect(),
                    ));
                } else if key == OUTPUT_NODE_TYPE {
                    values.push((
                        NodeType::Output,
                        value.iter().map(|op| op.0.clone()).collect(),
                    ));
                } else if key == VERTEX_NODE_TYPE {
                    values.push((
                        NodeType::Vertex,
                        value.iter().map(|op| op.0.clone()).collect(),
                    ));
                } else if key == EDGE_NODE_TYPE {
                    values.push((
                        NodeType::Edge,
                        value.iter().map(|op| op.0.clone()).collect(),
                    ));
                }
            }
        }

        let codec = match graph_type {
            Some("recurrent") => GraphCodec::recurrent(input_size, output_size, values),
            Some("weighted_directed") => {
                GraphCodec::weighted_directed(input_size, output_size, values)
            }
            Some("weighted_recurrent") => {
                GraphCodec::weighted_recurrent(input_size, output_size, values)
            }
            Some("gru") => GraphCodec::gru(input_size, output_size, values),
            Some("lstm") => GraphCodec::lstm(input_size, output_size, values),
            _ => GraphCodec::directed(input_size, output_size, values),
        };

        Ok(Self {
            codec: match max_nodes {
                Some(max_nodes) => codec.with_max_nodes(max_nodes),
                None => codec,
            },
        })
    }
}
