#!/usr/bin/env python

"""
A script to join some raw population and Deprivation stats together for each lsoa.
"""

import sys, argparse
from pathlib import Path
import smart_open
from codit import share_dir
import numpy as np
import pandas as pd

DATA_DIR = share_dir() / "codit" / "data" / "UK_regions"

POPULATION_XLS = DATA_DIR / "SAPE22DT13-mid-2019-lsoa-Broad_ages-estimates-unformatted.xlsx"
# www.ons.gov.uk/peoplepopulationandcommunity/populationandmigration/populationestimates/datasets/lowersuperoutputareamidyearpopulationestimatesnationalstatistics

DEPRIVATION_CSV = DATA_DIR / "imd2019lsoa.csv.gz"
# opendatacommunities.org/resource?uri=http%3A%2F%2Fopendatacommunities.org%2Fdata%2Fsocietal-wellbeing%2Fimd2019%2Findices

REGIONAL_CSV = DATA_DIR / "Output_Area_to_LSOA_to_MSOA_to_Local_Authority_District_(December_2017)_" \
                          "Lookup_with_Area_Classifications_in_Great_Britain.csv.gz"
# geoportal.statistics.gov.uk/datasets/fe6c55f0924b4734adf1cf7104a0173e_0/data

OUTPUT_CSV = share_dir() / "codit" / "data" / "city" / "population" / "sample_lsoa_population.csv.gz"

DEPRIVATION_COLUMNS = {
    ("a. Index of Multiple Deprivation (IMD)", "Rank"):                           "IMD_Rank",
    ("a. Index of Multiple Deprivation (IMD)", "Decile"):                         "IMD_Decile",
    ("a. Index of Multiple Deprivation (IMD)", "Score"):                          "IMD_Score",
    ("b. Income Deprivation Domain", "Rank"):                                     "Income_Rank",
    ("b. Income Deprivation Domain", "Decile"):                                   "Income_Decile",
    ("b. Income Deprivation Domain", "Score"):                                    "Income_Score",
    ("c. Employment Deprivation Domain", "Rank"):                                 "Employment_Rank",
    ("c. Employment Deprivation Domain", "Decile"):                               "Employment_Decile",
    ("c. Employment Deprivation Domain", "Score"):                                "Employment_Score",
    ("d. Education, Skills and Training Domain", "Rank"):                         "Education_Rank",
    ("d. Education, Skills and Training Domain", "Decile"):                       "Education_Decile",
    ("d. Education, Skills and Training Domain", "Score"):                        "Education_Score",
    ("e. Health Deprivation and Disability Domain", "Rank"):                      "Health_Rank",
    ("e. Health Deprivation and Disability Domain", "Decile"):                    "Health_Decile",
    ("e. Health Deprivation and Disability Domain", "Score"):                     "Health_Score",
    ("f. Crime Domain", "Rank"):                                                  "Crime_Rank",
    ("f. Crime Domain", "Decile"):                                                "Crime_Decile",
    ("f. Crime Domain", "Score"):                                                 "Crime_Score",
    ("g. Barriers to Housing and Services Domain", "Rank"):                       "Housing_Rank",
    ("g. Barriers to Housing and Services Domain", "Decile"):                     "Housing_Decile",
    ("g. Barriers to Housing and Services Domain", "Score"):                      "Housing_Score",
    ("h. Living Environment Deprivation Domain", "Rank"):                         "LivingEnv_Rank",
    ("h. Living Environment Deprivation Domain", "Decile"):                       "LivingEnv_Decile",
    ("h. Living Environment Deprivation Domain", "Score"):                        "LivingEnv_Score",
    ("i. Income Deprivation Affecting Children Index (IDACI)", "Rank"):           "IDACI_Rank",
    ("i. Income Deprivation Affecting Children Index (IDACI)", "Decile"):         "IDACI_Decile",
    ("i. Income Deprivation Affecting Children Index (IDACI)", "Score"):          "IDACI_Score",
    ("j. Income Deprivation Affecting Older People Index (IDAOPI)", "Rank"):      "IDAOPI_Rank",
    ("j. Income Deprivation Affecting Older People Index (IDAOPI)", "Decile"):    "IDAOPI_Decile",
    ("j. Income Deprivation Affecting Older People Index (IDAOPI)", "Score"):     "IDAOPI_Score",
}

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("population_xls", nargs="?", type=Path, default=POPULATION_XLS, help="Path to population excel file")
    parser.add_argument("deprivation_csv", nargs="?", type=Path, default=DEPRIVATION_CSV, help="Path to deprivation csv file")
    parser.add_argument("regional_csv", nargs="?", type=Path, default=REGIONAL_CSV, help="Path to regional lookup csv")
    parser.add_argument("output_csv", nargs="?", type=Path, default=OUTPUT_CSV, help="Path to population excel file")
    parser.add_argument("--sheet_name", default="Mid-2019 Persons", help="Sheet name for population xls file")
    parser.add_argument("-r", "--regions", nargs='+', help="Optional filter for regions to include (eg 'Leeds')")
    args = parser.parse_args()

    # Prevent national data overwriting the default region path
    if not args.regions and args.output_csv == OUTPUT_CSV:
        print("Please specify some regions when overwriting the default output_csv file", file=sys.stderr)
        parser.print_help()
        sys.exit(2)

    # Read CSV file and rename columns to match home_locations.py
    data = pd.read_excel(args.population_xls, args.sheet_name, index_col=0, usecols="A,B,G", skiprows=4)
    data.index.rename("lsoa11cd", inplace=True)
    data.rename(columns={"LSOA Name": "lsoa11nm", "All Ages": "population"}, inplace=True)

    # Filter out regions if specified
    if args.regions:
        mask = np.zeros(data.shape[0], dtype=np.bool_)
        for region in args.regions:
            mask |= data["lsoa11nm"].str.startswith(region)
        data = data[mask]

    # Load the regional_csv file:
    with smart_open.open(args.regional_csv) as fh:
        reg = pd.read_csv(fh)

    reg.columns = [c.lower() for c in reg.columns]
    reg_reduced = reg[['lsoa11cd', 'msoa11cd', 'msoa11nm', 'lad17cd', 'lad17nm']].drop_duplicates()
    # duplicates are due to this regional dataset being at an even finer granularity than lsoa
    reg_reduced.set_index('lsoa11cd', inplace=True)
    data = data.join(reg_reduced)

    # Load the deprivation_csv file
    with smart_open.open(args.deprivation_csv) as fh:
        dep = pd.read_csv(fh, index_col=0)

    # Join each deprivation index as a new column
    for dep_idx in sorted(dep["Indices of Deprivation"].unique()):
        for measure in dep["Measurement"].unique():
            where = (dep['Indices of Deprivation'] == dep_idx) & (dep['Measurement'] == measure)
            data = data.join(dep.loc[where]["Value"])
            column_name = DEPRIVATION_COLUMNS[(dep_idx.strip(), measure.strip())]
            data.rename(columns={"Value": column_name}, inplace=True)

    # Write output csv file
    with smart_open.open(args.output_csv, "w") as fh:
        data.to_csv(fh)

if __name__ == "__main__":
    main()
