#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
import sys, os
import numpy as np
import pickle

dirname = os.path.dirname(sys.argv[0])
if dirname:
    os.chdir(dirname)

fn_cache = 'data/supplement_b_cache.pickle'
if not os.path.exists(fn_cache):
    # To use this script, first visit
    # https://sedac.ciesin.columbia.edu/data/collection/gpw-v4
    # Create an account, and request data downloads for a single year from
    # https://sedac.ciesin.columbia.edu/data/set/gpw-v4-population-density-rev11/data-download
    # https://sedac.ciesin.columbia.edu/data/set/gpw-v4-population-density-rev11/data-download
    # https://sedac.ciesin.columbia.edu/data/set/gpw-v4-population-count-rev11/data-download
    # https://sedac.ciesin.columbia.edu/data/set/gpw-v4-national-identifier-grid-rev11/data-download
    # https://sedac.ciesin.columbia.edu/data/set/gpw-v4-land-water-area-rev11/data-download
    fn_density = 'data/gpw_v4_population_density_rev11_2020_2pt5_min.asc'
    fn_count = 'data/gpw_v4_population_count_rev11_2020_2pt5_min.asc'
    fn_admin = 'data/gpw_v4_national_identifier_grid_rev11_2pt5_min.asc'
    fn_land_area = 'data/gpw_v4_land_water_area_rev11_landareakm_2pt5_min.asc'
    def loadASCII(fn):
        with open(fn, 'r') as f:
            l=next(f); assert l.startswith('ncols'       ); ncols        = int  (l.rsplit(' ', 1)[1].strip())
            l=next(f); assert l.startswith('nrows'       ); nrows        = int  (l.rsplit(' ', 1)[1].strip())
            l=next(f); assert l.startswith('xllcorner'   ); xllcorner    = float(l.rsplit(' ', 1)[1].strip())
            l=next(f); assert l.startswith('yllcorner'   ); yllcorner    = float(l.rsplit(' ', 1)[1].strip())
            l=next(f); assert l.startswith('cellsize'    ); cellsize     = float(l.rsplit(' ', 1)[1].strip())
            l=next(f); assert l.startswith('NODATA_value'); NODATA_value = float(l.rsplit(' ', 1)[1].strip())
            table = np.loadtxt(f, float, ndmin=2)
            table = np.where(table == NODATA_value, np.nan, table)
            assert table.shape == (nrows, ncols)
            return ncols, nrows, xllcorner, yllcorner, cellsize, table
    # Load tables
    density = loadASCII(fn_density)
    count = loadASCII(fn_count)
    admin = loadASCII(fn_admin)
    landArea = loadASCII(fn_land_area)
    gridParams = density[:-1]
    assert count[:-1] == gridParams
    assert admin[:-1] == gridParams
    assert landArea[:-1] == gridParams
    cache = density, count, admin, landArea, gridParams
    pickle.dump(cache, open(fn_cache, 'wb'), -1)
else:
    cache = pickle.load(open(fn_cache, 'rb'))
    density, count, admin, landArea, gridParams = cache


density = density[-1]
count = count[-1]
admin = admin[-1]
landArea = landArea[-1]

# Service region population out of world population
serviceLatitudeRange = (-60., 60.) # [degrees]
ncols, nrows, xllcorner, yllcorner, cellsize = gridParams
serviceRegionPopFraction = np.nansum(
        count[
            int((serviceLatitudeRange[0]-yllcorner)/cellsize) :
            int((serviceLatitudeRange[1]-yllcorner)/cellsize)
        ]) / np.nansum(count)

# It has been assumed that Broadband LEO service will not be competitive with
# terrestrial broadband Internet in urban centers where such alternatives are
# available.

# Suppose that pmax is set according to the crude assumption that
# 42M Americans do not have the ability to purchase broadband Internet as of
# 2020, and that these people are all in low-density regions.
target_pop = 42e6
USA_cells = (admin==840).nonzero()
USA_density = density[USA_cells]
USA_count = count[USA_cells]
USA_cell_order = USA_density.argsort()
USA_pop = USA_count[USA_cell_order].cumsum()
lowDensityCells = USA_cell_order[:(USA_pop > target_pop).nonzero()[0][0]]
lowDensityIndices = USA_cells[0][lowDensityCells], USA_cells[1][lowDensityCells]
pmaxDeveloped = density[lowDensityIndices].max()

worldAverageFactor = 121.7 / 83.0
# We attempt to adjust pmax for differences outside the United States,
# referring to columns "World" and "Developed"
# in series "Active mobile-broadband subscriptions"
# in the figure at the bottom of page 5 of
# https://www.itu.int/en/ITU-D/Statistics/Documents/facts/FactsFigures2019.pdf
pmax = pmaxDeveloped * worldAverageFactor

# Now build a map of where potential customers are located.
lowDensityMask = np.zeros_like(count, dtype=bool)
lowDensityMask[density<pmax] = True
lowDensityCount = np.where(np.isnan(count), 0, count)
# How should we treat high-density areas?
if 0:
    # Cap them to a density of pmax
    lowDensityCount[~lowDensityMask] = pmax * landArea[~lowDensityMask]
else:
    # Zero them; these customers have other options
    lowDensityCount[~lowDensityMask] = 0

# For further processing, it will greatly speed things along to downsample at this point.
ncols, nrows, xllcorner, yllcorner, cellsize = gridParams
lowDensityCount = lowDensityCount.reshape(nrows//12, 12, ncols//12, 12).sum((1,3))
nrows//=12
ncols//=12
cellsize*=12

# For an NGSO system, one should not consider longitude.
lowDensityCount1D = lowDensityCount.sum(1)
lat = yllcorner + np.arange(nrows-.5, -.5, -1) * cellsize
lon = xllcorner + np.arange(ncols-.5, -.5, -1) * cellsize

# Over what range of latitudes is an SV visible?
rEarth = 6378e3 # [m]
altitude = 550e3 # [m]
minElevationAngle = 25 # [degrees]
serviceArcDegrees = 2 * np.rad2deg(np.arccos(rEarth/(rEarth+altitude) * np.cos(np.deg2rad(minElevationAngle))) - np.deg2rad(minElevationAngle))

# Suppose that a spherical cap of radius serviceArcDegrees/2 is located at
# latitude lat.  How many people live there?

# Given a unit vector n=[u,v,w], define the cap C by r.n ≥ d = cos(serviceArcDegrees/2) and integrate over longitude.
# ∫_C dΩ f
# = ∫_0^π sinθ dθ ∫_0^{2π} dφ 1_{r.n ≥ d}
# = ∫_0^π sinθ dθ ∫_0^{2π} dφ 1_{u cosφ sinθ + v sinφ sinθ + w cosθ ≥ d}
# Without loss of generality, let v = 0.
# = ∫_0^π sinθ dθ ∫_0^{2π} dφ 1_{u cosφ sinθ + w cosθ ≥ d}
# Now we have cases
# Case I:   (d - w cosθ) / (u sinθ) > +1
#   No integral
# Case II:  (d - w cosθ) / (u sinθ) < -1
#   Integrate over |φ| ≤ π
# Case III: |(d - w cosθ) / (u sinθ)| < 1
#   Integrate over |φ| ≤ arccos((d - w cosθ) / (u sinθ))
# The horizontal integral in each case is
# 2 arccos(clip((d - w cosθ) / (u sinθ), -1, +1))
# Expressed as a fraction of 2π, this is
# (1/π) arccos(clip((d - w cosθ) / (u sinθ), -1, +1))
# Transforming into latitude l = π/2-θ,
# (1/π) arccos(clip((d - w sin(l)) / (u cos(l)), -1, +1))
# Now write w = sin(l0), u = cos(l0).

# Compute the convolution of a spherical cap of radius serviceArcDegrees/2
# with the population count map.  Neither elegant nor fast.
d = np.cos(np.deg2rad(serviceArcDegrees/2))
clat = np.cos(np.deg2rad(lat))
slat = np.sin(np.deg2rad(lat))
integratedLowDensityCount = np.zeros_like(lowDensityCount)
lowDensityCountFT = np.fft.fft(lowDensityCount, axis=1)
for i,(u,w) in enumerate(zip(clat, slat)):
    # Build a 2-D mask for the spherical cap on the lat/lon grid.
    radius = (180/np.pi) * np.arccos(np.clip((d - w*slat) / (u*clat), -1, +1))
    weight = np.gradient(np.clip(lon[None,:], -radius[:,None], +radius[:,None]), axis=1) * -cellsize**-1
    # Convolve this mask by the count function in the east-west axis, and sum
    # in the north-south axis, to produce one row of output.
    weightFT = np.fft.fft(weight, axis=1)
    countFT = (lowDensityCountFT * weightFT).sum(0)
    integratedLowDensityCount[i] = np.fft.ifft(countFT).real

# Because our 2-D mask was centered around lon=0 in the middle of our array,
# the FFT-based convolution has introduced a rotation of 180°.  Undo it.
integratedLowDensityCount = np.fft.fftshift(integratedLowDensityCount, axes=1)

if 0:
    import pylab as pl
    pl.figure()
    ax = pl.subplot(1,2,1)
    ax.imshow(integratedLowDensityCount, extent=(xllcorner, xllcorner+ncols*cellsize, yllcorner, yllcorner+nrows*cellsize))
    ax = pl.subplot(1,2,2)
    ax.set_xlabel('Longitude (degrees)')
    ax.set_ylabel('Latitude (degrees)')
    ax.plot(integratedLowDensityCount.mean(1)/1e6, lat)
    ax.set_xlabel('Potential Customers (millions)')
    ax.set_ylabel('Latitude (degrees)')
    ax.set_title('Rural people in view of SV')

# Next we need an orbit to consider.  Sample various right ascensions for an
# inclination of 53 degrees, and sample the distribution of potential broadband
# LEO customers along the orbital track.
inclination = 53 # [degrees]
anomaly = np.arange(200)/200 * 2*np.pi
samples = []
for RA in np.arange(0, 360, 2):
    cRA, sRA = np.cos(np.deg2rad(RA)), np.sin(np.deg2rad(RA))
    cI, sI = np.cos(np.deg2rad(inclination)), np.sin(np.deg2rad(inclination))
    x0 = np.r_[+cRA,   +sRA,    0 ]
    x2 = np.r_[-sRA*sI,+cRA*sI, cI]
    x1 = np.cross(x2, x0)
    samplePoints = np.cos(anomaly)[:,None] * x0 + np.sin(anomaly)[:,None] * x1
    samplePointLat = np.arcsin(samplePoints[:,2])
    samplePointLon = np.arctan2(samplePoints[:,1], samplePoints[:,0])
    samplePointI = np.floor((np.rad2deg(samplePointLat)-yllcorner) / cellsize).astype(int)
    samplePointJ = np.floor((np.rad2deg(samplePointLon)-xllcorner) / cellsize).astype(int)
    samples.extend(integratedLowDensityCount[samplePointI, samplePointJ])

PARestimate = np.max(samples) / np.mean(samples)

from supplement_common import export

export('pmax', precision=1, siunit='\people\per\km^2')
export('pmaxDeveloped', precision=1, siunit='\people\per\km^2')
export('PARestimate', precision=1)
export('serviceRegionPopFraction', precision=1, scale=1e2, siunit='\percent')