# -------------------------------------------------------------------
# ORTHOGRAPHIC
# Your personal aerial satellite. Always on. At any altitude.*
# Developed by MarStrMind
# License: Open Software License 3.0
# Up to date version always on marstr.online
# -------------------------------------------------------------------
# xp_normalmap.py
# For some geographical features, we will add a normal map for
# additional realism in the orthos, without using additional meshes
# or geometry.
#
# Source converted to Python from this C++ StackOverflow:
# https://stackoverflow.com/a/2368794
# -------------------------------------------------------------------


import numpy as np
import math
import os
from PIL import Image, ImageFilter
from defines import *
from log import *


class mstr_xp_normalmap:

    # Only a few params
    def __init__(self, lat, lng, tag, value, tv, th, latlngfld):
        self._lat = lat
        self._lng = lng
        self._tag = tag
        self._value = value
        self._latlngfld = latlngfld
        self._tv = tv
        self._th = th
        mstr_msg("xp_normalmap", "[X-Plane] Normal Map generator initialized")


    # Load the layer image and resize it to 1/4th its size -
    # then provide it
    def load_layer(self):
        qtr = int(mstr_photores / 4)
        image = Image.open(mstr_datafolder + "_cache/" + str(self._lat) + "-" + str(self._tv) + "_" + str(self._lng) + "-" + str(self._th) + "_" + self._tag + "-" + self._value + "_layer.png")
        image = image.resize((qtr,qtr), Image.Resampling.LANCZOS)
        mstr_msg("xp_normalmap", "[X-Plane] Layer image loaded")
        return image


    # A few mathematical calls we need
    # --------------------------------------------------------
    def intensity(self, pixel):
        avg = (pixel[0] + pixel[1] + pixel[2]) / 3
        if avg > 0:
            pavg = 255.0 / avg
        else:
            pavg = 0
        return pavg


    def clamp(self, px, mpx):
        if px > mpx-1:
            return mpx-1
        else:
            if px < 0:
                return 0
            else:
                return px


    def map_component(self, px):
        return (px + 1.0) * (255.0 / 2.0)


    def normalize_vector(self, v):
        vc = np.array([v[0], v[1], v[2]])
        norm = np.linalg.norm(vc)
        nv = vc / norm
        return nv
    # --------------------------------------------------------


    # The Big Mac. Generate the normal map
    def generate_normal_map_for_layer(self, image):
        mstr_msg("xp_normalmap", "[X-Plane] Beginning normal map generation")
        # No specularity, no reflectivity - but standard color
        # Blue (reflectivity) and alpha (specularity) need to be 1 - but can be adjusted as needed

        # Resize original
        image = image.resize((int(mstr_photores/4), int(mstr_photores/4)), Image.Resampling.BILINEAR)

        nmp = Image.new("RGBA", (image.width, image.height), (128,128,1,1))
        org = image.load()
        nmp_pix = nmp.load()

        # Let's try some shenanigans
        w = image.width
        h = image.height
        for y in range(h):
            for x in range(w):
                p = org[x,y]
                if p[3] > 0: # Only do something if there is something to do in layer
                    # Neighboring pixels
                    px_t  = org[ self.clamp(x, w),   self.clamp(y+1, h) ]
                    px_tr = org[ self.clamp(x+1, w), self.clamp(y+1, h) ]
                    px_r  = org[ self.clamp(x+1, w), self.clamp(y, h)   ]
                    px_br = org[ self.clamp(x+1, w), self.clamp(y+1, h) ]
                    px_b  = org[ self.clamp(x, w),   self.clamp(y-1, h) ]
                    px_bl = org[ self.clamp(x-1, w), self.clamp(y-1, h) ]
                    px_l  = org[ self.clamp(x-1, w), self.clamp(y, h)   ]
                    px_tl = org[ self.clamp(x-1, w), self.clamp(y+1, h) ]

                    # Intensities of pixels
                    it_t  = self.intensity(px_t)
                    it_tr = self.intensity(px_tr)
                    it_r  = self.intensity(px_r)
                    it_br = self.intensity(px_br)
                    it_b  = self.intensity(px_b)
                    it_bl = self.intensity(px_bl)
                    it_l  = self.intensity(px_l)
                    it_tl = self.intensity(px_tl)

                    # Sobel filter
                    dx = (it_tr + 2.0 * it_r + it_br) - (it_tl + 2.0 * it_l + it_bl)
                    dy = (it_bl + 2.0 * it_b + it_br) - (it_tl + 2.0 * it_t + it_tr)
                    dz = 10 # This is usually a good value for strength
                    v = (dx, dy, dz)
                    nrm = self.normalize_vector(v)
                    
                    if nrm[1] > 0:
                        nrm[1] = 0 - (abs(nrm[1]))
                    else:
                        nrm[1] = abs(nrm[1])

                    # Set pixel
                    nmp_pix[x,y] = (int(self.map_component(nrm[0])), int(self.map_component(nrm[1])), 255 - int(self.map_component(nrm[2])), 1)

        mstr_msg("xp_normalmap", "[X-Plane] Normal map generated")
        return nmp
    

    # The funnction to call. Blends with the existing map, or creates a new one
    def build_normalmap(self, layer):
        mstr_msg("xp_normalmap", "[X-Plane] Building normal map")
        # The layer image
        #lyr = self.load_layer()

        # Make the normal map for the layer
        nrm = self.generate_normal_map_for_layer(layer)

        # Normal map final file name
        nrmfln = mstr_datafolder + "z_orthographic/normals/" + self._latlngfld + "/" + str(self._tv) + "_" + str(self._th) + ".png"

        # Check for existence of normal map file
        ex = os.path.isfile(nrmfln)

        # Does not exist? Just save
        if ex == False:
            nrm.save(nrmfln)

        # Exists? Open it, composite both, save
        if ex == True:
            nrmmap = Image.open(nrmfln)
            nrmmap.alpha_composite(nrm)

            # Specularity blending correction
            nrmmap_pix = nrmmap.load()
            for y in range(nrmmap.height):
                for x in range(nrmmap.width):
                    c = nrmmap_pix[x,y]
                    nrmmap_pix[x,y] = (c[0], c[1], c[2], 1)
            nrmmap.save(nrmfln)
            
        mstr_msg("xp_normalmap", "[X-Plane] Normal map saved")