#!/usr/bin/env python

import sys, os, shutil, re
import datetime as dt
import multiprocessing as mp
import subprocess
import atexit

gisbase = os.environ['GISBASE'] = r'C:\\Program Files\\GRASS GIS 7.2.0\\'
gisdbase = r'C:\data\grass'
sys.path.append(os.path.join(os.environ['GISBASE'], "etc", "python"))

import grass.script as grass
import grass.script.setup as gsetup

# code to create list with name of vectors to iterate over
className = "bykerne" + "Net"
mapName = className + '@Solar'          # counting layers (sample clusters)
result = grass.read_command('v.info',   # in vector file
                            flags = 'e',
                            map = mapName)

result = str.split(result, '\n')
layerCount = result[12]                 # dblink_num = layers 
layerCount = str.split(layerCount, '=')
layerCount = layerCount[1]              # number of layers
layerCount = int(layerCount)

layerN = 1                              # start layer count
outputName = className + '_' + str(layerN - 1)
vectorList = grass.read_command('g.list', type = 'vector', mapset = 'Solar')
vectorList = str.split(vectorList, '\n')
vectorList = ''.join(vectorList)
vectorList = str.split(vectorList, '\r')
vectorList.pop(0)
vectorList.pop(0)
vectorList.pop(-1) # final vector list

# parameters
location = 'Denmark'            
mhorizon = 'Solar'              
maxDistance = '500'             
hstep = 15                      
distance = '1.0'

# start of function definitions:

def worker_horizon(cpu,demh,azimuth, outputNameS):

        # Setup Grass Environment, including temporary mapset

        mtemp = 'temp'+str(cpu).zfill(2)

        gsetup.init(gisbase, gisdbase, location, mtemp)

        grass.run_command("g.region",vector = outputNameS +"@Solar")



        # map names

        prefhor = demh+'hor'

        demh = demh+'@Solar'

        

        # run r.horizon for one azimuth angle

        azi = str(azimuth)

        grass.run_command("r.horizon", elevation=demh, \

                          direction=azi, maxdistance=maxDistance, \

                          output=prefhor, dist=distance, overwrite = True)



        # rename horizon map to proper azimuth name

        az = str(azimuth).zfill(3)

        tempmap = prefhor+'_0'

        horizonmap = prefhor+'_'+az

        cmd = tempmap+","+horizonmap

        grass.run_command("g.rename",rast=cmd,overwrite = True)        



def create_temp(cores):

        gsetup.init(gisbase, gisdbase, location, 'PERMANENT')

        for count in range(0,cores,1):

                temp = 'temp'+str(count).zfill(2)

                temp_path = gisdbase+'/'+location+'/'+temp

                if(os.path.exists(temp_path) == False):

                        grass.run_command("g.mapset",flags="c", mapset=temp, quiet=1)

                        print(temp+" mapset created.")

                else:

                        print(temp+" mapset already exists. skipping...")



def remove_temp(cores):

        # setup is in the target mapset, not temp mapsets

        gsetup.init(gisbase, gisdbase, location, mhorizon)

        # Delete the temp mapset

        for count in range(0,cores):

                mapset_temp = 'temp'+str(count).zfill(2)

                grass.run_command("g.mapsets", mapset=mapset_temp, operation = 'remove')

                temp_path = gisdbase+'/'+location+'/'+mapset_temp

                shutil.rmtree(temp_path)
                
                if os.path.exists(temp_path):
                    shutil.rmtree(temp_path)


def copy_mapset(mapset_from,mapset_to,regfilter,overwrite):

        # list contents of temp mapset

        grass.run_command("g.mapset", mapset=mapset_from, quiet=1)

        raster_list = grass.list_pairs(type = 'rast')

        # Switch to target mapset and copy rasters over

        grass.run_command("g.mapset", mapset=mapset_to,quiet=1)

        for rast in raster_list:

                if(rast[1] != 'PERMANENT' and re.search(regfilter,rast[0])):

                        old = rast[0]+ '@' + rast[1]

                        new = rast[0] 

                        cmd = old+","+new

                        grass.run_command("g.copy", rast=cmd, overwrite= True)



def printout(str_text):

		timestamp = dt.datetime.strftime(dt.datetime.now(),"%H:%M:%S")

		lf.write(timestamp+": "+str_text+'\n')

		print timestamp+": "+str_text



def main():
        for i in vectorList:
            outputNameS = i
            dem = outputNameS
            
            ##################################

            # HORIZON CALCULATIONS

            ##################################

            global cores

            global lf

            cores = mp.cpu_count()



            # Open log file

            tlog = dt.datetime.strftime(dt.datetime.now(),"%Y-%m-%d_h%Hm%M")

            lf = open('rhorizon_'+tlog+'.log', 'w')



            # write params to log file

            print("STARTING R.HORIZON")


            # start logging        

            start = dt.datetime.now()

            starttime = dt.datetime.strftime(start,"%m-%d %H:%M:%S")

            print('START  r.horizon '+ starttime)



            # Create one temp directory for each CPU core

            print("Creating Temporary directories, one per cpu core.")

            create_temp(cores)



            # Spawn process for r.horizon DEM

            jobs = []

            c = 0
    
            for az in range(0,360,hstep):

                    if(c < cores):

                            p = mp.Process(target=worker_horizon, args=(c,dem,az, outputNameS))

                            p.start()

                            jobs.append(p)

                            pid = str(p.pid)

                            printout("R.Horizon: dem = "+dem+" process id = "+pid)

                            c += 1

                    else:	

                            # Wait for all the Processes to finish

                            print("R.Horizon. Cores maxed out.  Waiting for processes up to "+str(az)+" to finish.")

                            c = 0

                            for p in jobs:

                                    pid = str(p.pid)

                                    p.join()

                                    printout("R3. r.horizon. dem = "+dem+", pid = "+pid+" joined.")

            print("R.Horizon finished for "+dem)



            # Copy all the files back over to horizon mapset 

            print("R.Horizon finished. copying files back to "+mhorizon)

            for cpu in range(0,cores):

                    mtemp = 'temp'+str(cpu).zfill(2)

                    suf = 'hor_'

                    copy_mapset(mtemp, mhorizon, suf, 1)



            # Delete the temp mapsets
            print("Removing temporary mapsets")

            remove_temp(cores)



            # Stop Logging and record total processing time.

            end = dt.datetime.now()

            endtime = dt.datetime.strftime(end,"%m-%d %H:%M:%S")

            processingtime = end - start

            print('END r.horizon, processing time: '+str(processingtime))



            lf.close()
        
        sys.exit("Finished")

    
if __name__ == '__main__':
    #options, flags = grass.parser()
    #atexit.register(cleanup)
    main()