from PIL import Image
import copy
import numpy as np
import pickle
#-------Sample--------------------------------------(THIS WORKS)
# value = "01101000011001010110110001101100011011111"

# print(len(value))

# cmap = {'0': (0,155,3,0),
#         '1': (0,0,0)}

# data = [cmap[letter] for letter in value]
# img = Image.new('RGB', (8, (len(value)//8)+1), "white")
# print(data)
# img.putdata(data)
# img.show() 

#-------Sample---------------------------------------


class para_viz:
    def __init__(self,layers):
        self.layers=copy.deepcopy(layers)
        self.map()
    
    def map(self):
        for i in range(len(self.layers)):
            self.layers[i]=self.layers[i]/np.max(np.absolute(self.layers[i]),axis=0)
            self.layers[i]=self.layers[i]*255   
    
    def color_map(self,data):
        Data=[]
        for i in data:
            i=int(i)
            color=(i,0,0) if i>=0 else (0,i,0)
            Data.append(color)
        return Data

    def viz_neuron(self):
        neuron=self.layers[0][:,0]
        # print(neuron)
        flat_len=len(neuron)
        size=(flat_len//2,flat_len//2) if flat_len%2==0 else (flat_len//2,(flat_len//2)+1)
        img=Image.new('RGB',size,"white")
        data=self.color_map(neuron)
        print(len(data))
        img.putdata(data)
        img.show()


with open("weights.pkl","rb") as f:
    weights=pickle.load(f)

viz=para_viz(weights)
viz.viz_neuron()