text: '''
  Clean demographics data
  '''
  import os
  import numpy as np
  import pandas as pd
  import pickle
  
  CUR_PATH = os.path.split(os.path.abspath(__file__))[0]
  OUTPUT_DIR = os.path.join(CUR_PATH, 'data', 'demographic', 'clean')
  if os.access(OUTPUT_DIR, os.F_OK) is False:
      os.makedirs(OUTPUT_DIR)
  
  def select_year(data):
      '''
      Adjust data to be at the year-level
      '''
      data_adj = data.copy()
      data_adj = data_adj.pivot(index='Age', columns='Year', values='Values')
      return data_adj
  
  def select_fert_data(fert, set_zeroes=False):
      new_fert = fert[fert['AgeDef'] == 'ARDY']
      new_fert = new_fert[new_fert['Collection'] == 'HFD']
      new_fert = new_fert[(new_fert['RefCode'] == 'JPN_11')]
      new_fert.drop(['AgeDef', 'Collection', 'RefCode'], axis=1, inplace=True)
      new_fert.columns = ['Year', 'Age', 'Values']
      if set_zeroes:
          new_fert['Values'][new_fert['Age'] == 14] = 0
          new_fert['Values'][new_fert['Age'] == 15] = 0
          new_fert['Values'][new_fert['Age'] == 49] = 0
          new_fert['Values'][new_fert['Age'] == 50] = 0
      return new_fert.astype(float)
  
  def get_fert_data(fert_dir):
      fert_data = pd.read_csv(fert_dir, sep=r',\s*',\
          usecols=['Year1', 'Age', 'ASFR', 'AgeDef',\
                      'Collection', 'RefCode'], engine='python')
      fert_data = select_fert_data(fert_data)
      fert_data = select_year(fert_data)
  
      return fert_data
  
  def select_mort_pop_data(data):
      data_adj = data[data["Age"] != "110+"].astype(int)
      return data_adj
  
  def get_mort_pop_data(mort_dir, pop_dir):
      mort_data = pd.read_csv(mort_dir, sep=r'\s+', usecols=['Year', 'Age', 'Total'])
      mort_data = select_mort_pop_data(mort_data)
      pop_data = pd.read_csv(pop_dir, sep=r'\s+', usecols=['Year', 'Age', 'Total'])
      pop_data = select_mort_pop_data(pop_data)
  
      mort_data['Values'] = mort_data['Total'] / pop_data['Total']
      mort_data = mort_data.drop('Total', axis=1)
      mort_data = select_year(mort_data)
      mort_data = mort_data[mort_data.index < 100] # only concerned with ages 0-99
  
      pop_data['Values'] = pop_data['Total']
      pop_data = pop_data.drop('Total', axis=1)
      pop_data = select_year(pop_data)
      pop_data = pop_data[pop_data.index < 100]
  
      return mort_data, pop_data
  
  def calc_imm_resid(fert_t_minus_1, mort_t_minus_1, pop_t_minus_1, pop_t):
      '''
      Calculate immigration rate in year t
      as (pop_t - pop_{t-1}(1 - mort_{t-1})) / (pop_{t-1}),
      and setting pop_t_0 = pop_{t-1} * fert_{t-1}
      '''
      # First, calculate births
      if len(fert_t_minus_1) == 100:
          births = fert_t_minus_1 * pop_t_minus_1 / 2
      else:
          # Births from 14-50, but age start at 0
          births = fert_t_minus_1 * pop_t_minus_1[15:52] / 2
      births = np.sum(births)
      # Second, calculate deaths
      deaths = mort_t_minus_1 * pop_t_minus_1
      deaths = np.roll(deaths, 1)
      # Third, calculate predicted population
      pred_pop = np.roll(pop_t_minus_1, 1) - deaths
      pred_pop[0] = births
      # Fourth, calculate immigration rate
      imm = (pop_t - pred_pop) / (pop_t_minus_1)
      return imm
  
  def calc_imm_resids(fert_data, mort_data, pop_data):
      '''
      Calculate immigration rates over all data
      '''
      imm = pd.DataFrame()
      for year in range(1948, 2016):
          fert_t_minus_1 = fert_data[year - 1]
          mort_t_minus_1 = mort_data[year - 1]
          pop_t_minus_1 = pop_data[year - 1]
          pop_t = pop_data[year]
          imm[year] = calc_imm_resid(fert_t_minus_1, mort_t_minus_1, pop_t_minus_1, pop_t)
      return imm
  
  datadir = 'data/demographic/raw/'
  fert_dir = datadir + 'jpn_fertility.csv'
  mort_dir = datadir + 'jpn_mortality.csv'
  pop_dir = datadir + 'jpn_population.csv'
  
  # Extract fertility, mortality, and population data
  fert_data = get_fert_data(fert_dir)
  mort_data, pop_data = get_mort_pop_data(mort_dir, pop_dir)
  # Compute immigration rates
  imm_rate = calc_imm_resids(fert_data, mort_data, pop_data)
  # Keep last 3 years of immigration rates
  drop_immyears = sorted(imm_rate.columns)[:- 3]
  imm_rate = imm_rate.drop(drop_immyears, axis=1)
  
  # Save
  pickle.dump(fert_data, open(OUTPUT_DIR + 'fert.p', 'wb'))
  pickle.dump(mort_data, open(OUTPUT_DIR + 'mort.p', 'wb'))
  pickle.dump(pop_data, open(OUTPUT_DIR + 'pop.p', 'wb'))
  pickle.dump(imm_rate, open(OUTPUT_DIR + 'imm.p', 'wb'))
  