# -------------------------------------------------------------------
# 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
# -------------------------------------------------------------------
# maskgen.py
# The class that generates a mask of the layer it was asked to do.
# This mask will then be used to generate a photo layer, which in
# turn is then used to construct the final photo. It can be argued
# that this part of the code is the most crucial one, as the other
# classes involved rely on what this code is doing, and by extension,
# generating.
# -------------------------------------------------------------------

import math
from osmxml import *
from defines import *
from log import *
from PIL import Image, ImageFilter, ImageDraw, ImagePath
from random import randrange
from functions import *
import random

class mstr_maskgen:

    # Initializes the class with some required variables
    # Much of this code is adjusted to work within a class.
    def __init__(self, box, vstep, tag, value, isline, subtag=None, subvalue=None):
        self._box = box
        self._tag = tag
        self._subtag = subtag
        self._subvalue = subvalue
        self._value = value
        self._vstep = vstep
        self._scale = 1 / math.cos(math.radians(self._box[0]))
        self._isline = isline

        #mstr_msg("maskgen", "Intialized mask gen.")


    # Projects a point into the canvas of the mask.
    # Final projection depends on positive or negative latitude or longitude.
    def project_pixel(self, pnt, edge):
        pdiff = edge - pnt
        byT = pdiff * 1000
        divisor = byT / 16
        return divisor
    

    # Extract lat/lng from custom extracted nodes block
    def latlong_from_id(self, id, nds):
        latlng = []
        for i in nds:
            if i[0] == id:
                #latlng.append((float(i[1]), float(i[2])))
                latlng.append(float(i[1]))
                latlng.append(float(i[2]))
                break
        return latlng


    # Set width of tile - for buildings
    def set_tile_width(self, tile_width):
        self._tile_width = tile_width

    # Numbers needed for the possible building shadow layer
    def set_latlng_numbers(self, lat, tv, lng, th):
        self._latitude = lat
        self._lat_number = tv
        self._longitude = lng
        self._lng_number = th


    # Builds the required mask
    def _build_mask(self, xml, is_prep=False):
        # Generate empty image
        imgsize = 2048
        mask_img = Image.new("RGBA", (imgsize, imgsize))

        #tilexml = mstr_datafolder + "_cache/tile_" + str(self._latitude) + "-" + str(self._lat_number) + "_" + str(self._longitude) + "-" + str(self._lng_number) + ".xml"
        #tilexml = mstr_datafolder + "_cache/tile.xml"
        #xml = mstr_osmxml(0,0)
        fstr = str(self._box[0]) + "-" + str(self._box[1]) + "_" + str(self._box[2]) + "-" + str(self._box[3])
        nds = xml.acquire_nodes()
        way = xml.acquire_waypoint_data()
        rls = xml.acquire_relations()

        mstr_msg("maskgen", "Building mask for " + str(self._box[0]) + "-" + str(self._box[1]) + ", " + str(self._box[2]) + "-" + str(self._box[3]) + ", for " + self._tag + ": " + self._value )

        frs = []

        # Calculate actual bounding box
        bbox = []
        # Latitude
        bbox.append(self._box[0] + ((self._box[1]-1) * self._vstep))
        bbox.append(self._box[0] + ((self._box[1]-1) * self._vstep) + self._vstep)
        # Longitude
        bbox.append(self._box[2] + ((self._box[3]-1) * mstr_zl_18))
        bbox.append(self._box[2] + ((self._box[3]-1) * mstr_zl_18) + mstr_zl_18)

        # Building levels, if this is a building
        bld_levels = 0
        
        # Generate mask for ONE tag only
        if self._subtag == None:
            for w in way:
                if w[2] == self._tag and w[3] == self._value:
                    nd = []
                    for d in way:
                        if d[0] == w[0]:
                            if self._tag == "building" and bld_levels == 0:
                                bld_levels = xml.find_building_levels(w[0])
                            nd.append(d[1])
                    frs.append(nd)
            # Scout through relations as these also make up map data
            for r in rls:
                if self._tag in r[1] and self._value in r[1]:
                    nd = []
                    for w in way:
                        if int(w[0]) == int(r[0]):
                            nd.append(w[1])
                    frs.append(nd)
        
        # Generate mask for one tag, PLUS a subtag. This is mostly used for admin areas
        if self._subtag != None:
            nd = []
            wids = []
            for w in way:
                if w[2] == self._tag and w[3] == self._value:
                    wids.append(w[0])
            for w in wids:
                for wp in way:
                    if wp[0] == w and wp[2] == self._subtag and wp[3] in self._subvalue:
                        for d in way:
                            if d[0] == wp[0] and d[1] != "NULL":
                                nd.append(d[1])
                        frs.append(nd)

        # Project all pixels
        for f in frs:
            pts = []
            for a in f:
                latlng = self.latlong_from_id(a, nds)
                if len(latlng) == 2:
                    # For some reason, sometimes the array is empty. Make sure we have two data points.
                    if len(latlng) == 2:

                        # Project the pixel, and add to the polygon shape.
                        p_lat = self.project_pixel(latlng[0], bbox[1])
                        p_lng = self.project_pixel(latlng[1], bbox[3])
                        pixlat = 0
                        pixlng = 0
                        pr = 2048

                        # Draw pixels in direction according to latitude and longitude positions -

                        # Latitude:
                        if self._box[0] > 0:
                            pixlat = int((imgsize*self._scale)*p_lat)
                        if self._box[0] < 0:
                            pixlat = pr - (int((imgsize*self._scale)*p_lat))

                        # Longitude:
                        if self._box[2] > 0:
                            pixlng = int(imgsize - (imgsize*p_lng))
                        if self._box[2] < 0:
                            pixlng = pr - (int(imgsize - (imgsize*p_lng)))

                        pts.append((pixlng, pixlat))

            # Corel Draw!
            imgd = ImageDraw.Draw(mask_img)

            # Draw polygons
            if self._isline == False:
                if len(pts) >= 3:
                    if self._tag != "building":
                        imgd.polygon(pts, fill="#000000")
                    if self._tag == "building":
                        # Find ID of color index to use
                        idx = 0
                        for i in mstr_building_base_colors:
                            if i[0] == self._value:
                                break
                            else:
                                idx = idx + 1
                        # Now we have the index.
                        # Pick some color from it
                        c = randrange(len( mstr_building_base_colors[idx][1]))
                        clr = mstr_building_base_colors[idx][1][c]
                        # And draw the polygon with that -
                        # this will be the base color for that building in the layer
                        imgd.polygon(pts, fill=clr)

            # For road specific items, draw lines instead
            if self._isline == True:
                if len(pts) >= 2: # Only need two points to form a line
                    idx = 0
                    for i in range(len(mstr_ortho_layers)):
                        if mstr_ortho_layers[i][0] == self._tag and mstr_ortho_layers[i][1] == self._value:
                            idx = i
                            break
                    imgd.line(pts, fill="#000000", width=mstr_ortho_layers[idx][2], joint="curve")
        
        if is_prep == True:
            return mask_img


        # If this is a building, we need to render the shadow here, as we only know the height
        # of the building in this loop.
        if mstr_shadow_enabled == True and is_prep == False:
            if self._tag == "building":
                mpp = meters_per_pixel(self._tile_width) * mstr_zl_18
                pix_per_floor = mstr_shadow_floor_h / mpp
                total_pix = pix_per_floor * bld_levels
                shift = int(total_pix)

                fn = mstr_datafolder + "_cache/" + fstr + "_" + self._tag + "-" + self._value + "_layer_shadow.png"

                mask_pix = mask_img.load()

                bld_shadow = Image.new("RGBA", (mstr_photores, mstr_photores))
                bld_shadow_pix = bld_shadow.load()

                # Shadow sweep
                shf = 1
                while shf <= shift:
                    for y in range(mstr_photores):
                        for x in range(mstr_photores):
                            mp = mask_pix[x,y]
                            if mp[3] != 0:
                                if x+(shf*2) < mstr_photores and y+shf < mstr_photores:
                                    bld_shadow_pix[x+(shf*2), y+shf] = (0,0,0,255)
                    shf = shf+1

                # Building removal sweep
                for y in range(mstr_photores):
                    for x in range(mstr_photores):
                        mp = mask_pix[x,y]
                        if mp[3] != 0:
                            bld_shadow_pix[x,y] = (0,0,0,0)


                # Correct alpha
                bld_shadow_pix = bld_shadow.load()
                for y in range(mstr_photores):
                    for x in range(mstr_photores):
                        sp = bld_shadow_pix[x,y]
                        if sp[3] != 0:
                            bld_shadow_pix[x,y] = (0,0,0,120)

                # Store
                if os.path.isfile(fn) == True:
                    lyr = Image.open(fn)
                    lyr.alpha_composite(bld_shadow)
                    lyr.save(fn)
                else:
                    bld_shadow.save(fn)


        # Inform
        mstr_msg("maskgen", "Mask built.")

        # Return the image
        return mask_img