{# 
    Template file for generating the model specification. 
    Use the generate_models.py script to regenerate.
#}
// This file has been automatically generated by Jinja2 via the
// script {{ file }}.
// Please do not change this file by hand.
use blowtorch::ndarray::*;
use blowtorch::nn::loading::WeightLoader;
use blowtorch::nn::utils::Padding;
use blowtorch::nn::{Layer, FloatLikePrimitive};
use blowtorch::nn::{ConvolutionLayer, TransposedConvolutionLayer, LinearLayer, Flatten};


{% for m in models %} 
    pub struct {{m.module_name}}<F: FloatLikePrimitive> {
        {% for l in m.layers %}
            {{l.name}}: {{l.type_rust}}<F>,
        {% endfor %}
    }

    impl<F: FloatLikePrimitive> Layer<Array{{m.input_dim}}<F>, Array{{m.output_dim}}<F>> for {{m.module_name}}<F> {
        {# Have to allow since the last let might be extraneous due to model generation #}
        #[allow(clippy::let_and_return)]
        fn forward_pass(&self, input: &Array{{m.input_dim}}<F>) -> Array{{m.output_dim}}<F> {
            let x = input.clone();
            {% if debug %}
                trace!("input: {:?}\n", x);
            {% endif %}
            {% for l in m.layers %}
                let x = self.{{l.name}}.forward_pass(&x);
                {% if debug %}
                    trace!("{{m.python_name}}_{{l.python_name}}{{l.number}}_output: {:?}\n", x);
                {% endif %}
            {% endfor %}
            x
        }
    }

    impl<F: FloatLikePrimitive> {{m.module_name}}<F> {
        pub fn new(loader: &mut impl WeightLoader) -> Self {
            {% for l in m.layers -%}
                {% for w in l.weights -%}
                    {% if w is not none -%}
                        {% set weight_key = "layers" + "." + l.name + "." + w.name + ".npy" %}
                        let {{ l.name }}_{{w.name}} = loader.get_weight("{{weight_key}}",
                            {{ w.shape }}
                        ).unwrap();
                        {% if debug -%}
                            trace!("{{weight_key}}: {:?}\n", {{l.python_name}}{{l.number}}_{{w.name}});
                        {% endif -%}
                    {% endif -%}
                {% endfor -%}
                let {{l.name}} = {{l.type_rust}}::new(
                    
                    {% for w in l.weights -%}
                        {% if w is not none -%}
                            {% if w.optional -%}
                            Some({{ l.name }}_{{w.name}}),
                            {% else %}
                            {{ l.name }}_{{w.name}},
                            {% endif %}
                        {% else %}
                            None,
                        {% endif %}
                    {% endfor %}
                    {% for p in l.args_rust -%}
                    {{p}},
                    {%- endfor %}
                );
            {% endfor %}
            Self {
                {% for l in m.layers %}
                    {{l.name}},
                {% endfor %}
            }
        }
    }
{% endfor %}
