"""
Prepare ALI-Tax-Analyzer 2018 DIVA input file by reading
ALI-data/PIT_model_v4.1_inputdata.csv data, dropping unneeded
variables and renaming variables, and then writing the resulting
data/weights file to the ALI-Tax-Analyzer/alitaxanalyzer directory.

USAGE: execute this script in the ALI-Tax-Analyzer/data directory as follows:
$ python prep_diva18.py
"""

import os
import sys
import gzip
import shutil
import numpy as np
import pandas as pd
import taf  # Tax Analyzer Framework


TAX_YEAR = 2018

IN_DATA_PATH = os.path.join('..', '..', 'ALI-data',
                            'PIT_model_v4.1_inputdata.csv')
"""
% csvshow PIT_model_v4.1_inputdata.csv
1 TIN
2 She is a resident of Albania
3 He is not a resident of Albania
4 She is male
5 It is a woman
6 City/municipality/District
7 Marital status of tax filer
8 Status of children
9 Single tax filer
10 Statutory statement of the declareant(the deck learned is widowed)
11 Divorced
12 Status of the head of household/declarer
13 Key property ownership data
14 Key property ownership data for rent
15 Key property ownership data: free of charge
16 Key property ownership data: other
17 Other property
18 Next rental apartment
19 Next department: free of charge
20 Next department: other
21 Person under guardianship of the tax filer
22 Number of persons in custody of the family
23 Other persons under guardianship
24 (3) Gross income from wages or bonuses from employment relationships
25 (4) Estimated tax on salaries and rewards of employment
26 [5] Gross income from dividends
27 [6] Gross income from rent
28 (7) Capital gains from the sale of immovable property owned by the declarant
29 [8] Gross income from bank interest
30 (9) capital gains from real estate
31 (10) Shuma bruto të fituara nga lotari apo lojra të tjera të fatit
32 [11] Të ardhurat bruto nga pasuria intelektuale liçensa, ...
33 [12] Fitime kapitali nga dhurimi
34 [13] Të ardhura bruto të realizuara jashtë territorit të Republikës së ...
35 [14] Të ardhura bruto të tjera të papërmendura më lart
36 (15) SHUMA E TË ARDHURAVE BRUTO (kutia 3 + kutia 5 deri në 14)
37 (16) Shuma e interesit bankar të kredisë së marrë për shkollim
38 (17) Shpenzimet për mjekim për vete apo për fëmijët dhe për personat në ...
39 (18) SHUMA E SHPENZIMEVE TË ZBRITSHME (kutia 16 deri 17)
40 (19) Të ardhura të tatueshme pa të ardhura nga punësimi minus shpenzimet ...
41 (20) Tatimi i llogaritur për të ardhurat e tatueshme pa të ardhurat nga ...
42 (21) SHUMA E TATIMIT TË LLOGARITUR
43 (22) Tatimi mbi të ardhurat bruto nga pagat apo shpërblimet nga punësimi
44 (23) Tatimi mbi dividentin
45 (24) Tatimi mbi Qeranë
46 (25) Tatimi mbi fitimet kapitale nga shitja e pasurisë së paluajtshme
47 (26) Tatimi mbi interesat bankare
48 (27) Tatimi mbi fitimet kapitale të krijuara nga investimet në tituj ...
49 (28) Tatimi mbi fitimet nga llotaritë apo lojrat e tjera të fatit
50 (29) Tatimi mbi të ardhurat bruto nga të ardhurat sipas rubrikës
51 (30) Tatimi mbi të ardhurat nga fitimet kapitale nga dhurimi
52 (31) Tatimi mbi të ardhurat e tjera
53 (32) Tatimi mbi të ardhurat bruto të realizuara jashtë territorit ...
54 (33) Kreditimi i tatimit të huaj
55 (34) SHUMA E TATIMIT TË PAGUAR (kutia 22 deri kutia 32)
56 (35) TATIMI I DETYRUAR PËR T'U PAGUAR (kutia 21-kutia 34)
57 (36) TATIMI I PAGUAR TEPËR (kutia 34-kutia 21)
58 DRT
"""
IN_COLS = 58

NEW_VAR_NAMES = {         # [box] on the D1 Form
    'var1': 'recid',
    'var2': 'year',
    'var24': 'empinc',    # [3]
    'var25': 'emptaxwh',  # [4]
    'var26': 'divinc',    # [5]
    'var27': 'rentinc',   # [6]
    'var28': 'cgprop',    # [7]
    'var29': 'intinc',    # [8]
    'var30': 'cgfin',     # [9]
    'var31': 'gameinc',   # [10]
    'var32': 'ipinc',     # [11]
    'var33': 'cggifts',   # [12]
    'var34': 'forinc',    # [13]
    'var35': 'otherinc',  # [14]
    'var36': 'grossinc',  # [15]
    'var37': 'dedinted',  # [16]
    'var38': 'dedmedex',  # [17]
    'var39': 'dedtotal',  # [18]
    'var40': 'taxinc',    # [19]
    'var41': 'calctax',   # [20]
    'var42': 'pitax',     # [21]
}

DROP_VARS = [f'var{num}' for num in range(3, 23 + 1)] + \
            [f'var{num}' for num in range(43, IN_COLS + 1)]

EXPECT_FILEPATH = os.path.join('.', 'diva18_expect.csv')

CALC_VARS = ['emptaxwh', 'grossinc', 'dedtotal', 'taxinc', 'calctax', 'pitax']

OUT_DATA_PATH = os.path.join('..', 'alitaxanalyzer', 'diva18.csv')
WEIGHTS_PATH = os.path.join('..', 'alitaxanalyzer', 'diva18_weights.csv')


def compress(filename):
    """
    Compress specified file in gzip format with fixed file timestamp.
    """
    with open(filename, 'rb') as f_in:
        with gzip.GzipFile(filename=filename + '.gz',
                           mode='wb', mtime=0) as f_out:
            shutil.copyfileobj(f_in, f_out)


def main():
    """
    High-level logic of the script.
    """
    # read raw data file into a dataframe
    idf = pd.read_csv(IN_DATA_PATH,
                      # data contain Unicode
                      encoding='unicode_escape',
                      # some columns may have mixed data types
                      low_memory=False)
    # Side-step warning introduced when upgrading from pylint 2.7 to 2.9 in
    # which pylint does not see return value of pd.read_csv() as a dataframe:
    # pylint: disable=no-member
    numrecs = idf.shape[0]
    print(f'DIVA-{TAX_YEAR:d}:number_of_returns= {numrecs:d}')
    assert idf.shape[1] == IN_COLS

    # change variable names to var1, var2, ... , var58
    idf.columns = [f'var{num}' for num in range(1, IN_COLS + 1)]

    # drop unneeded variables
    idf.drop(columns=DROP_VARS, inplace=True)

    # rename variables
    idf.rename(columns=NEW_VAR_NAMES, inplace=True)

    # specify recid and year variables
    idf['recid'] = np.array(list(range(1, numrecs + 1)))
    idf['year'] = np.array([TAX_YEAR]*numrecs)

    # replace every blank field with a zero
    idf.fillna(0, inplace=True)

    # convert all variables to integer data type
    odf = idf.astype(np.int64)

    # write expect CSV file containing both model read and model calc variables
    taf.df2csv(odf, EXPECT_FILEPATH)
    # compress written CSV file in gzip format
    compress(EXPECT_FILEPATH)
    taf.delete_file(EXPECT_FILEPATH)

    # drop model calc variables
    odf.drop(columns=CALC_VARS, inplace=True)

    # write prepared dataframe to CSV file
    taf.df2csv(odf, OUT_DATA_PATH)

    # write companion weights to CSV file
    wght = pd.DataFrame()
    wght['WT2018'] = np.array([1]*numrecs, dtype=np.int32)
    taf.df2csv(wght, WEIGHTS_PATH)

    return 0


if __name__ == '__main__':
    sys.exit(main())
