#!/usr/bin/env python

import numpy as np

def print_matrix(m):
    acc = "["
    for i in range(m.shape[0]):
        row = "["
        for j in range(m.shape[1]):
            row += "{:.5f}".format(m[i, j])
            if j < m.shape[1]-1:
                row += ","
        row += "]"
        if i < m.shape[0]-1:
            row += ",\n"
        else:
            row += "\n"
        acc += row
    return acc + "]"


def print_matrix_as_row(m):
    m = m.flatten()
    acc = "["
    for i in range(m.shape[0]):
        acc += "{:.5f}".format(m[i])
        if i < m.shape[0]-1:
            acc += ","
        
    return acc + "]"

num_levels = 10

print("""
use ndarray::{Dim, ArrayBase, ArrayView, OwnedRepr};
use super::FeatureType;
"""
)

num_level = 256

I, J = np.ogrid[0:num_level, 0:num_level]
I = I.astype(np.float32)
J = J.astype(np.float32)


def calculate_weights(prop):
    if prop == "contrast":
        return (I-J) ** 2
    elif prop == "dissimilarity":
        return np.abs(I - J)
    elif prop == "homogeneity":
        return 1.0 / (1.0 + (I - J) ** 2)
    

print("const CONTRAST_WEIGHTS_{}: [f32; {}] = {};".format(
    num_level, 
    num_level * num_level, 
    print_matrix_as_row(calculate_weights("contrast"))
))
print("const DISSIMILARITY_WEIGHTS_{}: [f32; {}] = {};".format(
    num_level, 
    num_level * num_level, 
    print_matrix_as_row(calculate_weights("dissimilarity"))
))
print("const HOMOGENEITY_WEIGHTS_{}: [f32; {}] = {};".format(
    num_level, 
    num_level * num_level, 
    print_matrix_as_row(calculate_weights("homogeneity"))
))

print("""
pub fn get_weights(
    feature_type: FeatureType, 
    nlevels: usize) -> Option<ArrayBase<OwnedRepr<f32>, Dim<[usize; 4]>>>  {
    match (feature_type, nlevels) {
        (FeatureType::Contrast, 256) => 
            Some(ArrayView::from_shape(
                (256, 256, 1, 1), 
                &CONTRAST_WEIGHTS_256
            ).unwrap().to_owned()),
        (FeatureType::Dissimilarity, 256) => 
            Some(ArrayView::from_shape(
                (256, 256, 1, 1), 
                &DISSIMILARITY_WEIGHTS_256
            ).unwrap().to_owned()),
        (FeatureType::Homogeneity, 256) => 
            Some(ArrayView::from_shape(
                (256, 256, 1, 1), 
                &HOMOGENEITY_WEIGHTS_256
            ).unwrap().to_owned()),
        _ => None,
    }
}
"""
)

# for num_level in range(1, num_levels + 1):

#     I, J = np.ogrid[0:num_level, 0:num_level]
#     I = I.astype(np.float32)
#     J = J.astype(np.float32)

    
#     print("const WEIGHTS_{}: [[f32; {}]; {}] = {};".format(
#         num_level, 
#         num_level, 
#         num_level, 
#         print_matrix(I-J)))

# print()
# print("pub fn get_weights(nlevel: usize) -> Array2<f32> {")
# print("    match nlevel {")
# for i in range(1, num_levels + 1):
#     print("        {i} => ArrayView::from_shape(({i}, {i}), &WEIGHTS_{i}),".format(i=i))
# print("        _ => panic!(\"not supported\")")
# print("    }")
# print("}")