#!python
# encoding: utf-8
#
# Author: Klaus K. Holst
#
# targeted test program
#

import targeted as tg
import targeted.formula as tgf
import logging
import argparse
try:
    import coloredlogs
    coloredlogs.install(level='INFO')
except ModuleNotFoundError:
    logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

d = tg.getdata()

def test(n=1):
    if n == 0:
        print("Interaction term, correct nuisance model")
        val = tgf.riskreg(d, 'y~a', interaction='x', nuisance='x+z')
    elif n == 1:
        print("Correct nuisance model")
        val = tgf.riskreg(d, 'y~a', nuisance='x+z')
    elif n == 2:
        print("Incorrect nuisance model")
        val = tgf.riskreg(d, 'y~a', nuisance='x')
    else:
        print("Incorrect nuisance model")
        val = tgf.riskreg(d, 'y~a', nuisance='z')
    logger.info(val)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("n",
                        type=int,
                        nargs='?',  # optional argument
                        default=0,
                        help="type of model")
    parser.add_argument("-a", "--all",
                        action="store_true",
                        help="run all tests")

    args = parser.parse_args()
    print('targeted version ' + tg.__version__ + '\n')
    test(args.n)
