# -*- coding: utf-8 -*-
"""
/*------------------------------------------------------*
|                         BRAILS                        |
|                                                       |
| Author: Charles Wang,  UC Berkeley, c_w@berkeley.edu  |
|                                                       |
| Date:    01/5/2021                                    |
*------------------------------------------------------*/
"""

import argparse
import os
from pathlib import Path
import random
import numpy as np
from glob import glob
import pandas as pd
import geopandas as gpd
import json

'''
from brails.modules.RoofTypeClassifier import RoofClassifier
from .OccupancyClassClassifier import OccupancyClassifier
from brails.SoftstoryClassifier import SoftstoryClassifier
'''

from brails.modules import RoofClassifier, OccupancyClassifier, SoftstoryClassifier


from .workflow.Footprints import getMSFootprintsByPlace, getStateNameByBBox, getOSMFootprints, getMSFootprintsByBbox
from .workflow.Images import getGoogleImages



class CityBuilder:
    """Class for creating city-scale BIM."""

    def __init__(self, attributes=['story','occupancy','roofshape'], numBldg=10, random=True,bbox=[], place='',state='', footPrints='OSM', save=True, fileName='', workDir='tmp',GoogleMapAPIKey='', overwrite=False, reDownloadImgs=False):
        """init function for CityBuilder class.

        Args:
            place (str): Name of the city or county.
            state (str): Abbreviation of state name
            save (bool): Save file locally.
            fileName (str): Path of the file.

        Returns:
            BIM (geopandas dataframe): BIM
            fileName (str): Path to the footprint file.

        """

        if GoogleMapAPIKey == '':
            print('Please provide GoogleMapAPIKey') 
            return None

        if not os.path.exists(workDir): os.makedirs(workDir)

        #if fileName=='': fileName=Path(f'{place} {state}_BIM.geojson'.replace(' ','_').replace(',','_'))
        if fileName=='': 

            if len(bbox)>0: 
                fileName = Path(f'/{bbox[0]}_{bbox[1]}_{bbox[2]}_{bbox[3]}_BIM.geojson')
            else:
                fileName = Path(f'/{place}_BIM.geojson'.replace(' ','_').replace(',','_'))

        

        bimFile = Path(f'{workDir}/{fileName}')
        
        self.attributes = attributes
        self.numBldg = numBldg
        self.random = random
        self.place = place
        self.state = state
        self.save = save
        self.fileName = fileName
        self.workDir = workDir
        self.overwrite =  overwrite
        self.bimFile = bimFile
        self.GoogleMapAPIKey = GoogleMapAPIKey
        self.reDownloadImgs = reDownloadImgs
        self.footPrints = footPrints

        if os.path.exists(bimFile) and not overwrite:
            print('{} already exists.'.format(bimFile))
            BIM = gpd.read_file(bimFile)
            return None




        #BIM = getCityBIM()
        stateName = getStateNameByBBox(bbox=bbox,workDir=workDir,GoogleMapAPIKey=GoogleMapAPIKey) # isNone means didn't find in the U.S.
        if stateName:# bbox, in the US
            #self.BIM, footprintFilename = getMSFootprintsByBbox(bbox=bbox,stateName=stateName,save=save,workDir=workDir,overwrite=overwrite,GoogleMapAPIKey=GoogleMapAPIKey)
            if footPrints=='OSM':
                print(f"Trying to get OSM footprints by the bbox {bbox}.")
                self.BIM, footprintFilename = getOSMFootprints(bbox=bbox,save=save,workDir=workDir,overwrite=True)
            elif footPrints=='Microsoft':
                # MS
                print(f"Trying to get MS footprints by the bbox {bbox}.")
                self.BIM, footprintFilename = getMSFootprintsByBbox(bbox=bbox,stateName=stateName,save=save,workDir=workDir,overwrite=overwrite,GoogleMapAPIKey=GoogleMapAPIKey)
        else: # out of US or by place
            if len(bbox)==4: # out of US: use OSM no matter what
                print(f"Trying to get OSM footprints by the bbox {bbox}.")
                self.BIM, footprintFilename = getOSMFootprints(bbox=bbox,save=save,workDir=workDir,overwrite=True)
            else: # by place
                if footPrints=='OSM':
                    print(f"Trying to get OSM footprints by the place name {place}.")
                    self.BIM, footprintFilename = getOSMFootprints(place=place,save=save,workDir=workDir,overwrite=True)
                elif footPrints=='Microsoft':
                    assert False, f"Trying to get Microsoft footprints by the place name {place}. This only works for the US!"
                    self.BIM, footprintFilename = getMSFootprintsByPlace(place=place,state=state,save=save,workDir=workDir,overwrite=overwrite)
        
        '''
        else:
            if footPrints=='OSM':
                self.BIM, footprintFilename = getOSMFootprints(place=place,save=save,workDir=workDir,overwrite=overwrite)
            else:
                # This will cause error if the interested region is out of the US
                assert False, "Microsoft footprint only covers the United States."
                self.BIM, footprintFilename = getMSFootprintsByPlace(place=place,state=state,save=save,workDir=workDir,overwrite=overwrite)
        '''

        '''
        if self.BIM.shape[0] < 1:
            print('Didn\'t find footprints for this region.')
            return None
        '''
 
        if numBldg < self.BIM.shape[0]: 
            if random: self.BIM = self.BIM.sample(n=numBldg)
            else: self.BIM = self.BIM[:numBldg]
        


    def build(self):

        if os.path.exists(self.bimFile) and not self.overwrite:
            print('{} already exists.'.format(self.bimFile))
            BIM = gpd.read_file(self.bimFile)
            return BIM 
        
        if self.BIM.shape[0] < 1:
            print('Can not build because the Footprints module didn\'t find footprints for this region.')
            return None

        print('Starting downloading images.')
        imgDir = os.path.join(self.workDir, 'images')
        imageTypes = ['StreetView','TopView']
        getGoogleImages(self.BIM,GoogleMapAPIKey=self.GoogleMapAPIKey, imageTypes=imageTypes, imgDir=imgDir, ncpu=2,reDownloadImgs=self.reDownloadImgs)

        #BIM = allFootprints
        self.BIM['Lon'] = self.BIM['geometry'].centroid.x
        self.BIM['Lat'] = self.BIM['geometry'].centroid.y
        for cat in imageTypes:
            self.BIM[cat] = self.BIM.apply(lambda row: Path(f"{imgDir}/{cat}/{cat}x{row['Lon']}x{row['Lat']}.png"), axis=1)


        # initialize a roof classifier
        roofModel = RoofClassifier(printRes=False)

        # initialize an occupancy classifier
        occupancyModel = OccupancyClassifier(printRes=False)

        # initialize a soft-story classifier
        ssModel = SoftstoryClassifier(printRes=False)

        # use the roof classifier 
       
        roofShape_df = roofModel.predict(self.BIM['TopView'].tolist())

        # use the occupancy classifier 

        occupancy_df = occupancyModel.predict(self.BIM['StreetView'].tolist())

        # use the softstory classifier 

        softstory_df = ssModel.predict(self.BIM['StreetView'].tolist())

        self.BIM.reset_index(inplace = True) 
        self.BIM['ID'] = self.BIM.index

        

        roofShape = roofShape_df['prediction'].to_list()
        roofShapeProb = roofShape_df['probability'].to_list()
        self.BIM['roofShape'] = self.BIM.apply(lambda x: roofShape[x['ID']], axis=1)
        self.BIM['roofShapeProb'] = self.BIM.apply(lambda x: roofShapeProb[x['ID']], axis=1)

        softstory = softstory_df['prediction'].to_list()
        softstoryProb = softstory_df['probability'].to_list()
        self.BIM['softStory'] = self.BIM.apply(lambda x: softstory[x['ID']], axis=1)
        self.BIM['softStoryProb'] = self.BIM.apply(lambda x: roofShapeProb[x['ID']], axis=1)

        occupancy = occupancy_df['prediction'].to_list()
        occupancyProb = occupancy_df['probability'].to_list()
        self.BIM['occupancy'] = self.BIM.apply(lambda x: occupancy[x['ID']], axis=1)
        self.BIM['occupancyProb'] = self.BIM.apply(lambda x: occupancyProb[x['ID']], axis=1)

        # delete columns
        self.BIM.drop(columns=['Lat','Lon','StreetView','TopView','index'], axis=1, inplace=True)

        # save
        self.BIM.to_file(self.bimFile, driver='GeoJSON')
        print('BIM saved at {}'.format(self.bimFile))

        #print(self.BIM)
        return self.BIM








def main():
    pass

if __name__ == '__main__':
    main()




