import matplotlib.pyplot as plt
import numpy as np
from function import *

class Nnnar:
    def __init__(self, nbDimensions, minCoord, maxCoord, nbSubdivisions):
        self.nbDimensions = nbDimensions
        self.minCoord = minCoord
        self.maxCoord = maxCoord
        self.nbSubdivisions = nbSubdivisions
        self.delta = maxCoord - minCoord
        self.unit = self.delta / nbSubdivisions
        coord = np.zeros(nbDimensions,dtype=int) + self.nbSubdivisions
        self.space = np.zeros(coord)
        fillSpace = np.vectorize(lambda x: [], otypes=[object])
        self.space = fillSpace(self.space)
        self.calculatedSpaceAroundIdx = []
    
    def addPoint(self, coord, value):
        if len(coord) != self.nbDimensions:
            raise AttributeError("Error: wrong number of dimensions")
        self.space[*self.getSpaceIdxFromCoord(coord)].append(Point(coord, value))

    def getValueOfPoint(self, coord,nbNearest):
        points = self.getNNearest(coord, nbNearest)
        value = points[0].value
        for idx in range(1,len(points)):
            tvalue = points[idx].value
            for i in range(len(tvalue)):
                value[i] += tvalue[i]
        for i in range(len(value)):
            value[i] = value[i] / nbNearest
        return value

    def getNNearest(self, coord, nbNearest):
        idx = self.getSpaceIdxFromCoord(coord)
        nbAround = 0
        selected = []
        while len(selected) < nbNearest:
            selected = []
            subSpaceIdxList = getSpaceIdxAround(idx, nbAround, self)
            for subSpaceIdx in subSpaceIdxList:
                selected += selectPointsInRange(self.space[tuple(subSpaceIdx)],coord,nbAround*self.unit,self.space)
            nbAround += 1
        found = []
        dist = []
        i=0
        while len(found) < nbNearest:
            found.append(selected[i])
            dist.append(selected[i].getDistFromCoord(coord))
            i+=1
        maxDistIdx = foundMaxIndex(dist)
        maxDist = max(dist)
        for i in range(nbNearest, len(selected)):
            if (selected[i].getDistFromCoord(coord) < maxDist):
                found.pop(maxDistIdx)
                dist.pop(maxDistIdx)
                found.append(selected[i])
                dist.append(selected[i].getDistFromCoord(coord))
                maxDistIdx = foundMaxIndex(dist)
                maxDist = max(dist)
        return found

    def getSpaceIdxFromCoord(self,coord):
        coord = coord - self.minCoord
        coord = coord / self.unit
        coord = np.floor(coord).astype(int)
        return coord

class Point:
    def __init__(self, coord, value):
        self.coord = coord
        self.value = value

    def setvalue(self, value):
        self.value = value

    def getDist(self, point):
        return np.linalg.norm(self.coord - point.coord)

    def getDistFromCoord(self, coord):
        return np.linalg.norm(self.coord - coord)



n = Nnnar(3, 0, 3, 2)

n.addPoint(np.array([1, 0.8,1.5]), [0.2,0.5])
n.addPoint(np.array([0.2, 1.8,1.2]), [0.4,0.7])

print(n.getValueOfPoint(np.array([1, 1, 1]), 2))