
import random
import datetime
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed 
import warnings
import requests
from pathlib import Path
import geopandas as gpd
from tqdm import tqdm
from shapely.geometry import Point
import warnings
from shapely.errors import ShapelyDeprecationWarning
import math
warnings.filterwarnings("ignore", category=ShapelyDeprecationWarning)

import glob
import shutil
import numpy as np
import osmnx as ox
from typing import List, Union

from zensvi.download.base import BaseDownloader
import zensvi.download.mapillary.interface as mly
from zensvi.download.utils.imtool import ImageTool
from zensvi.download.utils.get_pids import panoids
from zensvi.download.utils.geoprocess import GeoProcessor
from zensvi.download.utils.helpers import standardize_column_names, create_buffer_gdf

class GSVDownloader(BaseDownloader):
    def __init__(self, gsv_api_key: str = None, log_path: str = None, distance: int = 1, grid: bool = False, grid_size: int = 1):
        """
        Google Street View Downloader class.

        Args:
            gsv_api_key (str, optional): Google Street View API key. Defaults to None.
            log_path (str, optional): Path to the log file. Defaults to None.
            distance (int, optional): Distance parameter for the GeoProcessor. Defaults to 1.
            grid (bool, optional): Grid parameter for the GeoProcessor. Defaults to False.
            grid_size (int, optional): Grid size parameter for the GeoProcessor. Defaults to 1.

        Raises:
            Warning: If gsv_api_key is not provided.
        """
        super().__init__(log_path, distance, grid, grid_size)
        if gsv_api_key == None:
            warnings.warn("Please provide your Google Street View API key to augment metadata.")
        self._gsv_api_key = gsv_api_key

    @property
    def gsv_api_key(self) -> str:
        """
        Property to get the Google Street View API key.

        Returns:
            str: Google Street View API key.
        """
        return self._gsv_api_key    
    
    @gsv_api_key.setter
    def gsv_api_key(self, gsv_api_key: str):
        """
        Setter to set the Google Street View API key.

        Args:
            gsv_api_key (str): Google Street View API key.
        """
        self._gsv_api_key = gsv_api_key

    def _augment_metadata(self, df):
        if self.cache_pids_augmented.exists():
            df = pd.read_csv(self.cache_pids_augmented)
            print("The augmented panorama IDs have been read from the cache")
            return df
        
        # Create a new directory called "augmented_metadata_checkpoints"
        dir_cache_augmented_metadata = self.dir_cache / 'augmented_pids'
        dir_cache_augmented_metadata.mkdir(parents=True, exist_ok=True)

        # Load all the checkpoint csv files
        checkpoints = glob.glob(str(dir_cache_augmented_metadata / '*.csv'))
        checkpoint_start_index = len(checkpoints)

        if checkpoint_start_index > 0:
            completed_rows = pd.concat([pd.read_csv(checkpoint) for checkpoint in checkpoints], ignore_index=True)
            completed_indices = completed_rows.index.unique()

            # Filter df to get remaining indices to augment metadata for
            df = df.loc[~df.index.isin(completed_indices)]

        def get_year_month(pid, proxies):
            url = "https://maps.googleapis.com/maps/api/streetview/metadata?pano={}&key={}".format(pid, self.gsv_api_key)
            while True:
                proxy = random.choice(proxies)
                try:
                    response = requests.get(url, proxies=proxy, timeout=5)
                    break
                except Exception as e:
                    print(f"Proxy {proxy} is not working. Exception: {e}")
                    continue

            response = response.json()
            if response['status'] == 'OK':
                # get year and month from date
                try:
                    date = response['date']
                    year = date.split("-")[0]
                    month = date.split("-")[1]
                except Exception:
                    year = None
                    month = None
                return {"year": year, "month": month}
            return {"year": None, "month": None}    

        def worker(row, proxies):
            panoid = row.panoid
            year_month = get_year_month(panoid, proxies)
            return row.Index, year_month

        batch_size = 1000  # Modify this to a suitable value
        num_batches = (len(df) + batch_size - 1) // batch_size

        for i in tqdm(range(num_batches), desc=f"Augmenting metadata by batch size {min(batch_size, len(df))}"):
            batch_df = df.iloc[i*batch_size : (i+1)*batch_size].copy()  # Copy the batch data to a new dataframe
            with ThreadPoolExecutor() as executor:
                batch_futures = {executor.submit(worker, row, self.proxies): row.Index for row in batch_df.itertuples()}
                for future in tqdm(as_completed(batch_futures), total=len(batch_futures), desc=f"Augmenting metadata for batch #{i+1}"):
                    row_index, year_month = future.result()
                    batch_df.at[row_index, 'year'] = year_month['year']
                    batch_df.at[row_index, 'month'] = year_month['month']

            # Save checkpoint for each batch
            batch_df.to_csv(f'{dir_cache_augmented_metadata}/checkpoint_batch_{checkpoint_start_index+i+1}.csv', index=False)
        
        # Merge all checkpoints into a single dataframe
        df = pd.concat([pd.read_csv(checkpoint) for checkpoint in glob.glob(str(dir_cache_augmented_metadata / '*.csv'))], ignore_index=True)

        # save the augmented metadata
        df.to_csv(self.cache_pids_augmented, index=False)
        # delete cache_lat_lon
        if self.cache_lat_lon.exists():
            self.cache_lat_lon.unlink() 
        # delete cache_pids_raw
        if self.cache_pids_raw.exists():
            self.cache_pids_raw.unlink()
        # delete the cache directory
        if dir_cache_augmented_metadata.exists():
            shutil.rmtree(dir_cache_augmented_metadata)
        return df

    def _get_pids_from_df(self, df, id_columns=None):
        # 1. Create a new directory called "pids" to store each batch pids
        dir_cache_pids = self.dir_cache / 'raw_pids'
        dir_cache_pids.mkdir(parents=True, exist_ok=True)

        # 2. Load all the checkpoint csv files
        checkpoints = glob.glob(str(dir_cache_pids / '*.csv'))
        checkpoint_start_index = len(checkpoints)
        if checkpoint_start_index > 0:
            dataframes = []
            for checkpoint in checkpoints:
                try:
                    df_checkpoint = pd.read_csv(checkpoint)
                    dataframes.append(df_checkpoint)
                except pd.errors.EmptyDataError:
                    print(f"Warning: {checkpoint} is empty and has been skipped.")
                    continue
            completed_rows = pd.concat(dataframes, ignore_index=True)

            completed_ids = completed_rows['lat_lon_id'].drop_duplicates()

            # Merge on the ID column, keeping track of where each row originates
            merged = df.merge(completed_ids, on='lat_lon_id', how='outer', indicator=True)

            # Filter out rows that come from the 'completed_ids' DataFrame
            df = merged[merged['_merge'] == 'left_only'].drop(columns='_merge')

        def get_street_view_info(longitude, latitude, proxies):
            results = panoids(latitude, longitude, proxies, )
            return results

        def worker(row):
            input_longitude = row.longitude
            input_latitude = row.latitude
            lat_lon_id = row.lat_lon_id
            id_dict = {column: getattr(row, column) for column in id_columns} if id_columns else {}
            return lat_lon_id, (input_longitude, input_latitude), get_street_view_info(input_longitude, input_latitude, self.proxies), id_dict

        # set lat_lon_id if it doesn't exist
        if 'lat_lon_id' not in df.columns:
            df['lat_lon_id'] = np.arange(1, len(df) + 1)
        results = []
        batch_size = 1000  # Modify this to a suitable value
        num_batches = (len(df) + batch_size - 1) // batch_size
        failed_rows = []
        
        # if there's no rows to process, return completed_ids
        if len(df) == 0:
            return completed_ids
        
        # if not, process the rows
        for i in tqdm(range(num_batches), desc=f"Getting pids by batch size {min(batch_size, len(df))}"):
            with ThreadPoolExecutor() as executor:
                batch_futures = {executor.submit(worker, row): row for row in df.iloc[i*batch_size : (i+1)*batch_size].itertuples()}
                for future in tqdm(as_completed(batch_futures), total=len(batch_futures), desc=f"Getting pids for batch #{i+1}"):
                    try:
                        lat_lon_id, (input_longitude, input_latitude), row_results, id_dict = future.result()
                        for result in row_results:
                            result['input_latitude'] = input_latitude
                            result['input_longitude'] = input_longitude
                            result['lat_lon_id'] = lat_lon_id
                            result.update(id_dict)
                            results.append(result)
                    except Exception as e:
                        print(f"Error: {e}")
                        failed_rows.append(batch_futures[future])  # Store the failed row

                # Save checkpoint for each batch
                if len(results) > 0:
                    pd.DataFrame(results).to_csv(f'{dir_cache_pids}/checkpoint_batch_{checkpoint_start_index+i+1}.csv', index=False)
                results = []  # Clear the results list for the next batch

        # Merge all checkpoints into a single dataframe
        results_df = pd.concat([pd.read_csv(checkpoint) for checkpoint in glob.glob(str(dir_cache_pids / '*.csv'))], ignore_index=True)

        # Retry failed rows
        if failed_rows:
            print("Retrying failed rows...")
            with ThreadPoolExecutor() as executor:
                retry_futures = {executor.submit(worker, row): row for row in failed_rows}
                for future in tqdm(as_completed(retry_futures), total=len(retry_futures), desc="Retrying failed rows"):
                    try:
                        lat_lon_id, (input_longitude, input_latitude), row_results, id_dict = future.result()
                        for result in row_results:
                            result['input_latitude'] = input_latitude
                            result['input_longitude'] = input_longitude
                            result['lat_lon_id'] = lat_lon_id
                            result.update(id_dict)
                            results.append(result)
                    except Exception as e:
                        print(f"Failed again: {e}")

            # Save the results of retried rows as another checkpoint
            if len(results) > 0:
                pd.DataFrame(results).to_csv(f'{dir_cache_pids}/checkpoint_retry.csv', index=False)
                # Merge the retry checkpoint into the final dataframe
                retry_df = pd.read_csv(f'{dir_cache_pids}/checkpoint_retry.csv')
                results_df = pd.concat([results_df, retry_df], ignore_index=True)

        # now save results_df as a new cache after dropping lat_lon_id
        results_df = results_df.drop(columns='lat_lon_id')
        # drop duplicates in panoid and id_columns
        results_df = results_df.drop_duplicates(subset=['panoid'] + id_columns)
        results_df.to_csv(self.cache_pids_raw, index=False)

        # delete the cache directory
        if dir_cache_pids.exists():
            shutil.rmtree(dir_cache_pids)
        return results_df

    def _get_pids_from_gdf(self, gdf, **kwargs):  
        if self.cache_lat_lon.exists():
            df = pd.read_csv(self.cache_lat_lon)
            print("The lat and lon have been read from the cache")
        else:
            if gdf.crs is None:
                gdf = gdf.set_crs('EPSG:4326')
            elif gdf.crs != 'EPSG:4326':
                # convert to EPSG:4326
                gdf = gdf.to_crs('EPSG:4326')
            # read shapefile
            gp = GeoProcessor(gdf, distance=self.distance, grid=self.grid, grid_size=self.grid_size, **kwargs)
            df = gp.get_lat_lon()
            df['lat_lon_id'] = np.arange(1, len(df) + 1)
            # save df to cache
            df.to_csv(self.cache_lat_lon, index=False)

        if self.cache_pids_raw.exists():
            print("The raw panorama IDs have been read from the cache")
            results_df = pd.read_csv(self.cache_pids_raw)
        else:
            # Use _get_pids_from_df to get pids from df
            results_df = self._get_pids_from_df(df, kwargs["id_columns"])

        # Check if lat and lon are within input polygons
        polygons = gpd.GeoSeries([geom for geom in gdf['geometry'] if geom.type in ['Polygon', 'MultiPolygon']])

        # the rest is only for polygons, so return results_df if there's no polygons
        if len(polygons) == 0:
            return results_df

        # Convert lat, lon to Points and create a GeoSeries
        points = gpd.GeoSeries([Point(lon, lat) for lon, lat in zip(results_df['lon'], results_df['lat'])])

        # Create a GeoDataFrame with the points and an index column
        points_gdf = gpd.GeoDataFrame(geometry=points, crs=gdf.crs)
        points_gdf['index'] = range(len(points_gdf))

        # Create a spatial index on the polygons GeoSeries
        polygons_sindex = polygons.sindex

        # Function to check whether a point is within any polygon
        def is_within_polygon(point):
            possible_matches_index = list(polygons_sindex.intersection(point.bounds))
            possible_matches = polygons.iloc[possible_matches_index]
            precise_matches = possible_matches.contains(point)
            return precise_matches.any()

        # Add progress bar for within_polygon calculation
        with tqdm(total=len(points), desc="Checking points within polygons") as pbar:
            within_polygon = []
            for point in points_gdf['geometry']:
                within_polygon.append(is_within_polygon(point))
                pbar.update()

        results_df['within_polygon'] = within_polygon

        # Return only those points within polygons
        results_within_polygons_df = results_df[results_df['within_polygon']]
        # Drop the 'within_polygon' column
        results_within_polygons_df = results_within_polygons_df.drop(columns='within_polygon')
        return results_within_polygons_df

    def _get_raw_pids(self, **kwargs):
        if self.cache_pids_raw.exists():
            pid = pd.read_csv(self.cache_pids_raw)
            print("The raw panorama IDs have been read from the cache")
            return pid

        if kwargs['lat'] is not None and kwargs['lon'] is not None:
            pid = panoids(kwargs['lat'], kwargs['lon'], self.proxies)
            pid = pd.DataFrame(pid)
            # add input_lat and input_lon
            pid['input_latitude'] = kwargs['lat']
            pid['input_longitude'] = kwargs['lon']
        elif kwargs['input_csv_file'] != "":
            df = pd.read_csv(kwargs['input_csv_file'])
            df = standardize_column_names(df)
            if kwargs['buffer'] > 0:
                gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.longitude, df.latitude), crs='EPSG:4326')
                gdf = create_buffer_gdf(gdf, kwargs['buffer'])
                pid = self._get_pids_from_gdf(gdf, **kwargs)
            else:
                pid = self._get_pids_from_df(df, kwargs['id_columns'])
        elif kwargs['input_shp_file'] != "":
            gdf = gpd.read_file(kwargs['input_shp_file'])
            if kwargs['buffer'] > 0:
                gdf = create_buffer_gdf(gdf, kwargs['buffer'])
            pid = self._get_pids_from_gdf(gdf, **kwargs)
        elif kwargs['input_place_name'] != "":
            print("Geocoding the input place name")
            gdf = ox.geocoder.geocode_to_gdf(kwargs['input_place_name'])
            # raise error if the input_place_name is not found
            if len(gdf) == 0:
                raise ValueError("The input_place_name is not found. Please try another place name.")
            if kwargs['buffer'] > 0:
                gdf = create_buffer_gdf(gdf, kwargs['buffer'])
            pid = self._get_pids_from_gdf(gdf, **kwargs) 
        else:
            raise ValueError("Please input the lat and lon, csv file, or shapefile.")

        return pid

    def _get_pids(self, path_pid, **kwargs):
        id_columns = kwargs['id_columns']
        if id_columns is not None:
            if isinstance(id_columns, str):
                id_columns = [id_columns.lower()]
            elif isinstance(id_columns, list):
                id_columns = [column.lower() for column in id_columns]
        else:
            id_columns = []
        # update id_columns
        kwargs['id_columns'] = id_columns
        
        # get raw pid
        pid = self._get_raw_pids(**kwargs)
        
        if kwargs["augment_metadata"] & (self.gsv_api_key != None):
            pid = self._augment_metadata(pid)
        elif kwargs["augment_metadata"] & (self.gsv_api_key == None):
            raise ValueError("Please set the gsv api key by calling the gsv_api_key method.")
        pid.to_csv(path_pid, index=False)
        print("The panorama IDs have been saved to {}".format(path_pid)) 

    def _set_dirs(self, dir_output):
        # set dir_output as attribute and create the directory
        self.dir_output = Path(dir_output)
        self.dir_output.mkdir(parents=True, exist_ok=True)
        # set dir_cache as attribute and create the directory
        self.dir_cache = self.dir_output / "cache_zensvi"
        self.dir_cache.mkdir(parents=True, exist_ok=True)
        # set other cache directories
        self.cache_lat_lon = self.dir_cache / "lat_lon.csv"
        self.cache_pids_raw = self.dir_cache / "pids_raw.csv"
        self.cache_pids_augmented = self.dir_cache / "pids_augemented.csv"
        
    def _filter_pids_date(self, pid_df, start_date, end_date):
        # create a temporary column date from year and month
        # Fill NA values with a placeholder value
        pid_df['year'] = pid_df['year'].fillna(-1)
        pid_df['month'] = pid_df['month'].fillna(-1)

        # Replace Inf values with a placeholder value
        pid_df['year'] = pid_df['year'].replace([np.inf, -np.inf], -1)
        pid_df['month'] = pid_df['month'].replace([np.inf, -np.inf], -1)

        # Convert to int and then to string, and create the 'date' column
        pid_df['date'] = pid_df['year'].astype(int).astype(str) + "-" + pid_df['month'].astype(int).astype(str)

        # Optionally, you can replace the placeholder values in 'date' column back to NaN
        pid_df['date'] = pid_df['date'].replace('-1--1', np.nan)

        # convert to datetime
        pid_df['date'] = pd.to_datetime(pid_df['date'], format="%Y-%m")
        # check if start_date and end_date are in the correct format with regex. If not, raise error
        if start_date is not None:
            try:
                start_date = datetime.datetime.strptime(start_date, "%Y-%m-%d")
            except ValueError:
                raise ValueError("Incorrect start_date format, should be YYYY-MM-DD")
        if end_date is not None:
            try:
                end_date = datetime.datetime.strptime(end_date, "%Y-%m-%d")
            except ValueError:
                raise ValueError("Incorrect end_date format, should be YYYY-MM-DD")
        # if start_date is not None, filter out the rows with date < start_date
        pid_df = pid_df[pid_df['date'] >= start_date] if start_date is not None else pid_df
        # if end_date is not None, filter out the rows with date > end_date
        pid_df = pid_df[pid_df['date'] <= end_date] if end_date is not None else pid_df
        # drop the temporary column date
        pid_df = pid_df.drop(columns='date')
        return pid_df
    
    def download_svi(self, dir_output: str, path_pid: str = None, zoom: int = 2, 
                      cropped: bool = False, full: bool = True, lat: float = None, lon: float = None, 
                      input_csv_file: str = "", input_shp_file: str = "", input_place_name: str = "", 
                      id_columns: Union[str, List[str]] = None, buffer: int = 0, augment_metadata: bool = False, 
                      batch_size: int = 1000, update_pids: bool = False, start_date: str = None, end_date: str = None, 
                      metadata_only: bool = False, **kwargs) -> None:
        """
        Downloads street view images.

        Args:
            dir_output (str): The output directory.
            path_pid (str, optional): The path to the panorama ID file. Defaults to None.
            zoom (int, optional): The zoom level for the images. Defaults to 2.
            cropped (bool, optional): Whether to crop the images. Defaults to False.
            full (bool, optional): Whether to download full images. Defaults to True.
            lat (float, optional): The latitude for the images. Defaults to None.
            lon (float, optional): The longitude for the images. Defaults to None.
            input_csv_file (str, optional): The input CSV file. Defaults to "".
            input_shp_file (str, optional): The input shapefile. Defaults to "".
            input_place_name (str, optional): The input place name. Defaults to "".
            id_columns (Union[str, List[str]], optional): The ID columns. Defaults to None.
            buffer (int, optional): The buffer size. Defaults to 0.
            augment_metadata (bool, optional): Whether to augment the metadata. Defaults to False.
            batch_size (int, optional): The batch size for downloading. Defaults to 1000.
            update_pids (bool, optional): Whether to update the panorama IDs. Defaults to False.
            start_date (str, optional): The start date for the panorama IDs. Format is isoformat (YYYY-MM-DD). Defaults to None.
            end_date (str, optional): The end date for the panorama IDs. Format is isoformat (YYYY-MM-DD). Defaults to None.
            metadata_only (bool, optional): Whether to download metadata only. Defaults to False.
            **kwargs: Additional keyword arguments.

        Returns:
            None
        """
        # set necessary directories
        self._set_dirs(dir_output)
        
        # call _get_pids function first if path_pid is None
        if (path_pid is None) & (self.cache_pids_augmented.exists() == False):
            print("Getting pids...")
            path_pid = self.dir_output / "gsv_pids.csv"
            if path_pid.exists() & (update_pids == False):
                print("update_pids is set to False. So the following csv file will be used: {}".format(path_pid))
            else:
                self._get_pids(path_pid, lat=lat, lon=lon,
                            input_csv_file=input_csv_file, input_shp_file = input_shp_file, input_place_name = input_place_name, 
                            id_columns=id_columns, buffer = buffer, augment_metadata=augment_metadata, **kwargs)
        elif self.cache_pids_augmented.exists():
            # copy the cache pids_augmented to path_pid
            path_pid = self.dir_output / "gsv_pids.csv"
            shutil.copy2(self.cache_pids_augmented, path_pid)
            print("The augmented panorama IDs have been saved to {}".format(path_pid))
        # Horizontal Google Street View tiles
        # zoom 3: (8, 4); zoom 5: (26, 13) zoom 2: (4, 2) zoom 1: (2, 1);4:(8,16)
        # zoom = 2
        # h_tiles = 4  # 26
        # v_tiles = 2  # 13
        # cropped = False
        # full = True
        # stop if metadata_only is True
        if metadata_only:
            print("The metadata has been downloaded")
            return
        # create a folder within self.dir_output
        self.panorama_output = self.dir_output / "gsv_panorama"
        self.panorama_output.mkdir(parents=True, exist_ok=True)
        
        panoids = self._read_pids(path_pid, start_date, end_date)
        
        if len(panoids) == 0:
            print("There is no panorama ID to download")
            return
        else:
            panoids_rest = self._check_already(panoids)

        if len(panoids_rest) > 0:
            UAs = random.choices(self.user_agents, k = len(panoids_rest))
            # check zoom level is 0<=zoom<=6
            if zoom < 0 or zoom > 6:
                raise ValueError("zoom level should be between 0 and 6")
            h_tiles = 2 ** zoom
            v_tiles = math.ceil(h_tiles / 2)
            ImageTool.dwl_multiple(panoids_rest, zoom, v_tiles, h_tiles, self.panorama_output, UAs, self.proxies, cropped, full, batch_size = batch_size, log_path=self.log_path)
        else:
            print("All images have been downloaded")
        
        # delete the cache directory
        if self.dir_cache.exists():
            shutil.rmtree(self.dir_cache)
            print("The cache directory has been deleted")