<reponame>igor-morawski/reverse-undistortion
import os 
import os.path as op
import json
import numpy as np
import copy
import pickle

class Aligner:
    def __init__(self, H, W, window_size : int, stride, addresses_LUT = None, mappings = None, raw_preprocessing_transforms = None, log=False, load_from=None):
        ''' Maps are numbered in C-order (like np.flatten)'''
        self._req_attributes = ["H", "W", "window_size", "stride", "addresses_LUT", "mappings", "_rptc", "_maps_n"]
        self._att_to_serialize = {"addresses_LUT" : int, "mappings" : None, "_rptc" : None}
        assert isinstance(window_size, int)
        assert stride <= window_size
        assert not H % window_size 
        assert not W % window_size
        assert not H % stride 
        assert not W % stride

        self.log = log
        if load_from:
            _H = H
            _W = W
            _window_size = window_size
            _stride = stride
            self.load(load_from)
            b = [self.H != _H, self.W != W, self.window_size != _window_size, self.stride != _stride]
            if any(b):
                self._log_msg("Overriding aligner init.", log=self.log)
                print(self.H, _H, ";", self.W, W, ";", self.window_size, _window_size, ";", self.stride, _stride)
            return

        if not isinstance(raw_preprocessing_transforms, type(None)):
            #raw_preprocessing_transforms composed
            if isinstance(raw_preprocessing_transforms, Transforms):
                self._rptc = np.linalg.inv(raw_preprocessing_transforms.compose())
            elif isinstance(raw_preprocessing_transforms, np.ndarray):
                assert raw_preprocessing_transforms.shape == (3,3)
                self._rptc =  np.linalg.inv(raw_preprocessing_transforms)
        else: 
            self._rptc = np.eye(3)

        self.H = H
        self.W = W
        self.window_size = window_size
        self.stride = stride
    
        self._maps_n = (self._apply_conv_formula(self.H + self.window_size, self.window_size, padding=0, stride=self.stride, assert_divisible=True)) * \
            (self._apply_conv_formula(self.W + self.window_size, self.window_size, padding=0, stride=self.stride, assert_divisible=True))

        if isinstance(addresses_LUT, type(None)):
            print("[WARNING] Constructing LUTs is costly; cache if possible.")
            self.addresses_LUT = self._construct_LUT_adresses()
        else:
            assert addresses_LUT.shape == (self.H, self.W)
            self.addresses_LUT = addresses_LUT
        assert self.addresses_LUT.max() + 1 == self._maps_n #
        if not mappings:
            self.mappings = {i : None for i in range(self._maps_n)}
        assert max(self.mappings.keys()) == self._maps_n - 1
        return
        
    
    def _apply_conv_formula(self, input, kernel, padding, stride, assert_divisible=False):
        result = int(int(input-kernel+2*padding)/stride)
        if assert_divisible:
            result == (input-kernel+2*padding)/stride
        return result

    def loop_over_locations(self, fnc, kwargs):
        for y in range(0, self.H - self.window_size + self.window_size, self.stride):
            for x in range(0, self.W - self.window_size + self.window_size, self.stride):
                fnc(x=x, y=y, **kwargs)

    def _construct_LUT_adresses(self):
        distances = float("inf")*np.ones((self.H, self.W))
        addresses = -1*np.ones((self.H, self.W), dtype=int)
        ys = -1*np.ones((self.H, self.W))
        xs = -1*np.ones((self.H, self.W))
        for i in range(self.H):
            for j in range(self.W):
                ys[i, j] = i
                xs[i, j] = j 
        idx = 0
        for y in range(0, self.H - self.window_size + self.window_size, self.stride):
            for x in range(0, self.W - self.window_size + self.window_size, self.stride):
                distances_window = distances[y:y + self.window_size, x:x + self.window_size]
                if not distances_window.shape == (self.window_size, self.window_size):
                    dws = distances_window.shape
                    ws = self.window_size
                    assert (dws[0] == ws and dws[1] == ws//2) or (dws[1] == ws and dws[0] == ws//2)  or (dws[1] == ws//2 and dws[0] == ws//2)
                dist_y = ys[y:y + self.window_size, x:x + self.window_size] - y + self.window_size//2
                dist_x = xs[y:y + self.window_size, x:x + self.window_size] - x + self.window_size//2
                dist = np.sqrt(dist_y**2 + dist_x**2)
                indices = np.where(dist < distances_window)
                distances[indices[0]+y, indices[1]+x] = dist[indices]
                addresses[indices[0]+y, indices[1]+x] = idx
                idx += 1
        assert (addresses >= 0).all()
        return addresses

        
    def apply(self, x, y, floor=True) -> list:
        if x==self.W:
            x-=1
        if y==self.H:
            y-=1
        mapping_idx = self.addresses_LUT[y, x] 
        mapping = self.mappings[mapping_idx]
        _result = mapping @ np.array([x, y, 1])
        _result = self._rptc @ _result
        if floor:
            _result = _result.astype(int)
        return (_result[0], _result[1])

    def save(self, filepath):
        data = {}
        for attribute in self._req_attributes:
            data[attribute] = getattr(self, attribute)
        np.save(filepath, data)
        return True

    def load(self, filepath):
        data = np.load(filepath, allow_pickle=True).item()
        for req_key in self._req_attributes:
            if req_key not in data.keys(): 
                raise ValueError(f"Required key {req_key} not found in {filepath}!")
            setattr(self, req_key, data[req_key])
        self._log_msg(f"Loaded {filepath}", log=self.log)
        return True
        
    def _save_json(self, filepath):
        data = {}
        for attribute in self._req_attributes:
            data[attribute] = getattr(self, attribute)
            if attribute in self._att_to_serialize.keys():
                if isinstance(data[attribute], np.ndarray):
                    serialized = np.array(data[attribute]).copy().tolist()
                elif isinstance(data[attribute], dict):
                    serialized = {}
                    for k, v in copy.deepcopy(data[attribute]).items():
                        a = data[attribute][k]
                        serialized_a = np.array(a).tolist()
                        serialized[int(k)] = serialized_a
                else: 
                    raise ValueError
                data[attribute] = serialized
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=4)
        self._log_msg(f"Aligner saved to {filepath}", log=self.log)
        return True


    def _load_json(self, filepath):
        # XXX needs a fix << self.mappings.keys() are strings instead of integers!
        with open(filepath) as f:
            data = json.load(f)
        for req_key in self._req_attributes:
            if req_key not in data.keys(): 
                raise ValueError(f"Required key {req_key} not found in {filepath}!")
            setattr(self, req_key, data[req_key])
            if req_key in self._att_to_serialize.keys():
                 data[req_key] = np.array(data[req_key], dtype=self._att_to_serialize[req_key])
            #     data[req_key] = data[req_key].reshape([3,3], dtype=self._att_to_serialize[req_key])
        self._log_msg(f"Loaded {filepath}", log=self.log)
        return True
        
        

    def _log_msg(self, msg, log=True):
        print(msg)
        if log:
            with open('align.log', 'a+') as f:
                f.write(msg+"\n")
        return True


class Transforms:
    def __init__(self):
        self.transforms = []
        self.descriptions = []
        self.jpeg_transforms = []
        self.raw_transforms = []
        self.jpeg_descriptions = []
        self.raw_descriptions = []
    
    def log(self, mtx, img, desc=None):
        trnsfm = mtx.copy()
        assert trnsfm.shape == (2,2) or trnsfm.shape == (3,3) 
        if trnsfm.shape == (2,2):
            trnsfm = np.eye(3)
            trnsfm[0:2,0:2] = mtx
        if not desc:
            desc = ""
        else:
            desc = "_" + desc
        desc = img + desc
        if img == "jpeg":
            self.jpeg_transforms.append(trnsfm)
            self.jpeg_descriptions.append(desc)
        elif img == "raw":
            self.raw_transforms.append(trnsfm)
            self.raw_descriptions.append(desc)
        else:
            raise NotImplementedError
        self._update()
    
    def _update(self):
        self.transforms = self._raw_operator(self.raw_transforms) + \
             self._jpeg_operator(self.jpeg_transforms) 
        self.descriptions = self._raw_operator(self.raw_descriptions) + \
             self._jpeg_operator(self.jpeg_descriptions) 
    
    def compose(self):
        t = np.eye(3)
        for m in self.transforms:
            t = m @ t
        return t 
    def _is_l_of_a(self, l):
        for e in l:
            if not isinstance(e, np.ndarray): return False
        return True

    def _jpeg_operator(self, l):
        return l.copy()[::-1]
        
    def _raw_operator(self, l):
        if self._is_l_of_a(l):
            return [np.linalg.inv(a) for a in l]
        else: 
            return l.copy()

    def print(self):
        for idx, (trnsfm, desc) in enumerate(zip(self.transforms, self.descriptions)):
            print(f"{idx}. {desc}: ")
            print(f"{trnsfm}")

class Transform_Fnc:
    def __init__(self, mtx, floor=True) -> None:
        assert mtx.shape == (3,3)
        self.mtx = mtx
        self.floor = floor
    
    def apply(self, x, y) -> list:
        _result = self.mtx @ np.array([x, y, 1])
        if self.floor:
            _result = _result.astype(int)
        return (_result[0], _result[1])

import os 
import os.path as op
import json
import numpy as np
import copy
import pickle

class Aligner:
    def __init__(self, H, W, window_size : int, stride, addresses_LUT = None, mappings = None, raw_preprocessing_transforms = None, log=False, load_from=None):
        ''' Maps are numbered in C-order (like np.flatten)'''
        self._req_attributes = ["H", "W", "window_size", "stride", "addresses_LUT", "mappings", "_rptc", "_maps_n"]
        self._att_to_serialize = {"addresses_LUT" : int, "mappings" : None, "_rptc" : None}
        assert isinstance(window_size, int)
        assert stride <= window_size
        assert not H % window_size 
        assert not W % window_size
        assert not H % stride 
        assert not W % stride

        self.log = log
        if load_from:
            _H = H
            _W = W
            _window_size = window_size
            _stride = stride
            self.load(load_from)
            b = [self.H != _H, self.W != W, self.window_size != _window_size, self.stride != _stride]
            if any(b):
                self._log_msg("Overriding aligner init.", log=self.log)
                print(self.H, _H, ";", self.W, W, ";", self.window_size, _window_size, ";", self.stride, _stride)
            return

        if not isinstance(raw_preprocessing_transforms, type(None)):
            #raw_preprocessing_transforms composed
            if isinstance(raw_preprocessing_transforms, Transforms):
                self._rptc = np.linalg.inv(raw_preprocessing_transforms.compose())
            elif isinstance(raw_preprocessing_transforms, np.ndarray):
                assert raw_preprocessing_transforms.shape == (3,3)
                self._rptc =  np.linalg.inv(raw_preprocessing_transforms)
        else: 
            self._rptc = np.eye(3)

        self.H = H
        self.W = W
        self.window_size = window_size
        self.stride = stride
    
        self._maps_n = (self._apply_conv_formula(self.H + self.window_size, self.window_size, padding=0, stride=self.stride, assert_divisible=True)) * \
            (self._apply_conv_formula(self.W + self.window_size, self.window_size, padding=0, stride=self.stride, assert_divisible=True))

        if isinstance(addresses_LUT, type(None)):
            print("[WARNING] Constructing LUTs is costly; cache if possible.")
            self.addresses_LUT = self._construct_LUT_adresses()
        else:
            assert addresses_LUT.shape == (self.H, self.W)
            self.addresses_LUT = addresses_LUT
        assert self.addresses_LUT.max() + 1 == self._maps_n #
        if not mappings:
            self.mappings = {i : None for i in range(self._maps_n)}
        assert max(self.mappings.keys()) == self._maps_n - 1
        return
        
    
    def _apply_conv_formula(self, input, kernel, padding, stride, assert_divisible=False):
        result = int(int(input-kernel+2*padding)/stride)
        if assert_divisible:
            result == (input-kernel+2*padding)/stride
        return result

    def loop_over_locations(self, fnc, kwargs):
        for y in range(0, self.H - self.window_size + self.window_size, self.stride):
            for x in range(0, self.W - self.window_size + self.window_size, self.stride):
                fnc(x=x, y=y, **kwargs)

    def _construct_LUT_adresses(self):
        distances = float("inf")*np.ones((self.H, self.W))
        addresses = -1*np.ones((self.H, self.W), dtype=int)
        ys = -1*np.ones((self.H, self.W))
        xs = -1*np.ones((self.H, self.W))
        for i in range(self.H):
            for j in range(self.W):
                ys[i, j] = i
                xs[i, j] = j 
        idx = 0
        for y in range(0, self.H - self.window_size + self.window_size, self.stride):
            for x in range(0, self.W - self.window_size + self.window_size, self.stride):
                distances_window = distances[y:y + self.window_size, x:x + self.window_size]
                if not distances_window.shape == (self.window_size, self.window_size):
                    dws = distances_window.shape
                    ws = self.window_size
                    assert (dws[0] == ws and dws[1] == ws//2) or (dws[1] == ws and dws[0] == ws//2)  or (dws[1] == ws//2 and dws[0] == ws//2)
                dist_y = ys[y:y + self.window_size, x:x + self.window_size] - y + self.window_size//2
                dist_x = xs[y:y + self.window_size, x:x + self.window_size] - x + self.window_size//2
                dist = np.sqrt(dist_y**2 + dist_x**2)
                indices = np.where(dist < distances_window)
                distances[indices[0]+y, indices[1]+x] = dist[indices]
                addresses[indices[0]+y, indices[1]+x] = idx
                idx += 1
        assert (addresses >= 0).all()
        return addresses

        
    def apply(self, x, y, floor=True) -> list:
        if x==self.W:
            x-=1
        if y==self.H:
            y-=1
        mapping_idx = self.addresses_LUT[y, x] 
        mapping = self.mappings[mapping_idx]
        _result = mapping @ np.array([x, y, 1])
        _result = self._rptc @ _result
        if floor:
            _result = _result.astype(int)
        return (_result[0], _result[1])

    def save(self, filepath):
        data = {}
        for attribute in self._req_attributes:
            data[attribute] = getattr(self, attribute)
        np.save(filepath, data)
        return True

    def load(self, filepath):
        data = np.load(filepath, allow_pickle=True).item()
        for req_key in self._req_attributes:
            if req_key not in data.keys(): 
                raise ValueError(f"Required key {req_key} not found in {filepath}!")
            setattr(self, req_key, data[req_key])
        self._log_msg(f"Loaded {filepath}", log=self.log)
        return True
        
    def _save_json(self, filepath):
        data = {}
        for attribute in self._req_attributes:
            data[attribute] = getattr(self, attribute)
            if attribute in self._att_to_serialize.keys():
                if isinstance(data[attribute], np.ndarray):
                    serialized = np.array(data[attribute]).copy().tolist()
                elif isinstance(data[attribute], dict):
                    serialized = {}
                    for k, v in copy.deepcopy(data[attribute]).items():
                        a = data[attribute][k]
                        serialized_a = np.array(a).tolist()
                        serialized[int(k)] = serialized_a
                else: 
                    raise ValueError
                data[attribute] = serialized
        with open(filepath, 'w') as f:
            json.dump(data, f, indent=4)
        self._log_msg(f"Aligner saved to {filepath}", log=self.log)
        return True


    def _load_json(self, filepath):
        # XXX needs a fix << self.mappings.keys() are strings instead of integers!
        with open(filepath) as f:
            data = json.load(f)
        for req_key in self._req_attributes:
            if req_key not in data.keys(): 
                raise ValueError(f"Required key {req_key} not found in {filepath}!")
            setattr(self, req_key, data[req_key])
            if req_key in self._att_to_serialize.keys():
                 data[req_key] = np.array(data[req_key], dtype=self._att_to_serialize[req_key])
            #     data[req_key] = data[req_key].reshape([3,3], dtype=self._att_to_serialize[req_key])
        self._log_msg(f"Loaded {filepath}", log=self.log)
        return True
        
        

    def _log_msg(self, msg, log=True):
        print(msg)
        if log:
            with open('align.log', 'a+') as f:
                f.write(msg+"\n")
        return True


class Transforms:
    def __init__(self):
        self.transforms = []
        self.descriptions = []
        self.jpeg_transforms = []
        self.raw_transforms = []
        self.jpeg_descriptions = []
        self.raw_descriptions = []
    
    def log(self, mtx, img, desc=None):
        trnsfm = mtx.copy()
        assert trnsfm.shape == (2,2) or trnsfm.shape == (3,3) 
        if trnsfm.shape == (2,2):
            trnsfm = np.eye(3)
            trnsfm[0:2,0:2] = mtx
        if not desc:
            desc = ""
        else:
            desc = "_" + desc
        desc = img + desc
        if img == "jpeg":
            self.jpeg_transforms.append(trnsfm)
            self.jpeg_descriptions.append(desc)
        elif img == "raw":
            self.raw_transforms.append(trnsfm)
            self.raw_descriptions.append(desc)
        else:
            raise NotImplementedError
        self._update()
    
    def _update(self):
        self.transforms = self._raw_operator(self.raw_transforms) + \
             self._jpeg_operator(self.jpeg_transforms) 
        self.descriptions = self._raw_operator(self.raw_descriptions) + \
             self._jpeg_operator(self.jpeg_descriptions) 
    
    def compose(self):
        t = np.eye(3)
        for m in self.transforms:
            t = m @ t
        return t 
    def _is_l_of_a(self, l):
        for e in l:
            if not isinstance(e, np.ndarray): return False
        return True

    def _jpeg_operator(self, l):
        return l.copy()[::-1]
        
    def _raw_operator(self, l):
        if self._is_l_of_a(l):
            return [np.linalg.inv(a) for a in l]
        else: 
            return l.copy()

    def print(self):
        for idx, (trnsfm, desc) in enumerate(zip(self.transforms, self.descriptions)):
            print(f"{idx}. {desc}: ")
            print(f"{trnsfm}")

class Transform_Fnc:
    def __init__(self, mtx, floor=True) -> None:
        assert mtx.shape == (3,3)
        self.mtx = mtx
        self.floor = floor
    
    def apply(self, x, y) -> list:
        _result = self.mtx @ np.array([x, y, 1])
        if self.floor:
            _result = _result.astype(int)
        return (_result[0], _result[1])
Maps are numbered in C-order (like np.flatten) #raw_preprocessing_transforms composed # # XXX needs a fix << self.mappings.keys() are strings instead of integers! #     data[req_key] = data[req_key].reshape([3,3], dtype=self._att_to_serialize[req_key])