import json
import re

from awacs.aws import PolicyDocument, Statement, Allow, Principal
from awacs.ecr import GetAuthorizationToken, BatchCheckLayerAvailability, GetDownloadUrlForLayer, BatchGetImage
from awacs.logs import CreateLogStream, PutLogEvents
from awacs.secretsmanager import GetSecretValue
from awacs.sts import AssumeRole
from cfn_flip import to_yaml
from stringcase import pascalcase
from troposphere import GetAtt, Output, Parameter, Ref, Sub, Split, Select
from troposphere.applicationautoscaling import ScalableTarget
from troposphere.applicationautoscaling import ScalingPolicy, TargetTrackingScalingPolicyConfiguration, \
    PredefinedMetricSpecification, CustomizedMetricSpecification, MetricDimension as AppAutoscalingMetricDimension
from troposphere.cloudwatch import Alarm, MetricDimension
from troposphere.ec2 import SecurityGroup
from troposphere.ecs import (AwsvpcConfiguration, DeploymentConfiguration, LoadBalancer, NetworkConfiguration,
                             PlacementStrategy,
                             Service, Tags)
from troposphere.elasticloadbalancingv2 import (Action, Certificate, Listener, ListenerRule, Condition,
                                                HostHeaderConfig, PathPatternConfig)
from troposphere.elasticloadbalancingv2 import LoadBalancer as ALBLoadBalancer
from troposphere.elasticloadbalancingv2 import LoadBalancer as NLBLoadBalancer
from troposphere.elasticloadbalancingv2 import (Matcher, RedirectConfig,
                                                TargetGroup,
                                                TargetGroupAttribute)
from troposphere.elasticloadbalancingv2 import SubnetMapping
from troposphere.iam import Role, Policy

from cloudlift.config import DecimalEncoder
from cloudlift.config import get_account_id
from cloudlift.config.region import get_environment_level_alb_listener
from cloudlift.config.service_configuration import DEFAULT_TARGET_GROUP_DEREGISTRATION_DELAY, \
    DEFAULT_LOAD_BALANCING_ALGORITHM, DEFAULT_HEALTH_CHECK_INTERVAL_SECONDS, DEFAULT_HEALTH_CHECK_TIMEOUT_SECONDS, \
    DEFAULT_HEALTH_CHECK_HEALTHY_THRESHOLD_COUNT, DEFAULT_HEALTH_CHECK_UNHEALTHY_THRESHOLD_COUNT
from cloudlift.deployment.deployer import build_config, get_automated_injected_secret_name
from cloudlift.deployment.launch_types import LAUNCH_TYPE_FARGATE, LAUNCH_TYPE_EC2, get_launch_type
from cloudlift.deployment.task_definition_builder import TaskDefinitionBuilder, container_name
from cloudlift.deployment.template_generator import TemplateGenerator
from cloudlift.exceptions import UnrecoverableException

class ServiceTemplateGenerator(TemplateGenerator):
    PLACEMENT_STRATEGIES = [
        PlacementStrategy(
            Type='spread',
            Field='attribute:ecs.availability-zone'
        )]

    def __init__(self, service_configuration, environment_stack, env_sample_file, ecr_image_uri, access_file,
                 desired_counts=None):
        super(ServiceTemplateGenerator, self).__init__(service_configuration.environment)
        self._derive_configuration(service_configuration)
        self.env_sample_file_path = env_sample_file
        self.environment_stack = environment_stack
        self.ecr_image_uri = ecr_image_uri
        self.access_file = access_file
        self.desired_counts = desired_counts or {}

    def _derive_configuration(self, service_configuration):
        self.application_name = service_configuration.service_name
        self.configuration = service_configuration.get_config()

    def generate_service(self):
        self._add_service_parameters()
        self._add_service_outputs()
        self.ecs_service_role = self._add_ecs_service_iam_role()
        self._add_cluster_services()
        self._add_ecr_outputs()
        return to_yaml(self.template.to_json())

    def _add_cluster_services(self):
        for ecs_service_name, config in self.configuration['services'].items():
            self._add_service(ecs_service_name, config)

    def _add_service_alarms(self, svc, config):
        # How to add service task count alarm
        # http://docs.aws.amazon.com/AmazonECS/latest/developerguide/cloudwatch-metrics.html#cw_running_task_count
        ecs_no_running_tasks_alarm = Alarm(
            'EcsNoRunningTasksAlarm' + str(svc.name),
            EvaluationPeriods=int(config.get('ecs_no_running_tasks_evaluation_periods', 2)),
            Dimensions=[
                MetricDimension(
                    Name='ClusterName',
                    Value=self.cluster_name
                ),
                MetricDimension(
                    Name='ServiceName',
                    Value=GetAtt(svc, 'Name')
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription=self._enrich_alarm_description('Alarm if the task count goes to zero, denoting service is down', config),
            Namespace='AWS/ECS',
            Period=int(config.get('ecs_no_running_tasks_period', 300)),
            ComparisonOperator='LessThanThreshold',
            Statistic='SampleCount',
            Threshold='1',
            MetricName='CPUUtilization',
            TreatMissingData='breaching'
        )
        self.template.add_resource(ecs_no_running_tasks_alarm)

    def _add_scalable_target(self, ecs_svc, config):
        resource_id = Sub('service/' + self.cluster_name + '/' + '${service_name}',
                          service_name=GetAtt(ecs_svc, "Name"))
        scalable_target = ScalableTarget(
            str(ecs_svc.name) + "ScalableTarget",
            MinCapacity=int(config.get('min_capacity')),
            MaxCapacity=int(config.get('max_capacity')),
            ResourceId=resource_id,
            RoleARN=f'arn:aws:iam::{self.account_id}:role/aws-service-role/ecs.application-autoscaling.amazonaws.com/AWSServiceRoleForApplicationAutoScaling_ECSService',
            ScalableDimension='ecs:service:DesiredCount',
            ServiceNamespace='ecs'
        )
        self.template.add_resource(scalable_target)
        return scalable_target

    def _add_scalable_target_alarms(self, service_name, ecs_svc, config):
        max_scalable_target_alarm = Alarm(
            'MaxScalableTargetAlarm' + service_name,
            EvaluationPeriods=3,
            DatapointsToAlarm=3,
            Dimensions=[
                MetricDimension(
                    Name='ServiceName',
                    Value=GetAtt(ecs_svc, 'Name')
                ),
                MetricDimension(
                    Name='ClusterName',
                    Value=self.cluster_name
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription=self._enrich_alarm_description('Triggers if desired task count of a service is equal to max_capacity, review auto scaling configuration if this alarm triggers', config),
            Namespace='ECS/ContainerInsights',
            Period=300,
            ComparisonOperator='GreaterThanOrEqualToThreshold',
            Statistic='Maximum',
            Threshold=int(config.get('max_capacity')),
            MetricName='DesiredTaskCount',
            TreatMissingData='notBreaching'
        )
        self.template.add_resource(max_scalable_target_alarm)

    def _add_alb_request_count_scaling_policy(self, ecs_svc, alb_arn, target_group, config, scalable_target):
        try:
            target_value = int(config.get('target_value'))
            scale_in_cool_down = int(config.get('scale_in_cool_down_seconds'))
            scale_out_cool_down = int(config.get('scale_out_cool_down_seconds'))
        except TypeError as e:
            raise UnrecoverableException('The following value has to be integer: {}'.format(e))

        if type(alb_arn) == str:
            alb_name = alb_arn.split('/')[2]
            alb_id = alb_arn.split('/')[3]
        else:
            alb_name = Select(2, Split('/', Ref(alb_arn)))
            alb_id = Select(3, Split('/', Ref(alb_arn)))

        tg_name = Select(1, Split('/', Ref(target_group)))
        tg_id = Select(2, Split('/', Ref(target_group)))
        self.template.add_resource(
            ScalingPolicy(
                str(ecs_svc.name) + 'ALBRequestCountPerTargetScalingPolicy',
                PolicyName='requestCountPerTarget',
                PolicyType='TargetTrackingScaling',
                TargetTrackingScalingPolicyConfiguration=TargetTrackingScalingPolicyConfiguration(
                    ScaleInCooldown=scale_in_cool_down,
                    ScaleOutCooldown=scale_out_cool_down,
                    TargetValue=target_value,
                    PredefinedMetricSpecification=PredefinedMetricSpecification(
                        PredefinedMetricType='ALBRequestCountPerTarget',
                        ResourceLabel=Sub("app/${alb_name}/${alb_id}/targetgroup/${tg_name}/${tg_id}", alb_id=alb_id,
                                          alb_name=alb_name, tg_name=tg_name, tg_id=tg_id)
                    )
                ),
                ScalingTargetId=Ref(scalable_target)
            )
        )

    def _add_custom_metric_scaling_policy(self, ecs_svc, config, scalable_target):
        try:
            target_value = int(config.get('target_value'))
            scale_in_cool_down = int(config.get('scale_in_cool_down_seconds'))
            scale_out_cool_down = int(config.get('scale_out_cool_down_seconds'))
        except TypeError as e:
            raise UnrecoverableException('The following value has to be integer: {}'.format(e))

        metric_dimensions = [AppAutoscalingMetricDimension(
            Name=d['name'], Value=d['value'],
        ) for d in config.get('metric_dimensions')]

        self.template.add_resource(
            ScalingPolicy(
                str(ecs_svc.name) + 'CustomMetricScalingPolicy',
                PolicyName='customMetricScalingPolicy',
                PolicyType='TargetTrackingScaling',
                TargetTrackingScalingPolicyConfiguration=TargetTrackingScalingPolicyConfiguration(
                    ScaleInCooldown=scale_in_cool_down,
                    ScaleOutCooldown=scale_out_cool_down,
                    TargetValue=target_value,
                    CustomizedMetricSpecification=CustomizedMetricSpecification(
                        MetricName=config.get('metric_name'),
                        Namespace=config.get('namespace'),
                        Statistic=config.get('statistic'),
                        Unit=config.get('unit'),
                        Dimensions=metric_dimensions,
                    ),

                ),
                ScalingTargetId=Ref(scalable_target)
            )
        )

    def _add_service(self, service_name, config):
        launch_type = get_launch_type(config)
        secrets_name = config.get('secrets_name')
        container_configurations = build_config(
            env_name=self.env,
            service_name=self.application_name,
            ecs_service_name=service_name,
            sample_env_file_path=self.env_sample_file_path,
            essential_container_name=container_name(service_name),
            secrets_name=secrets_name,
            access_role=config.get('access_role'),
            access_file=self.access_file,
            override_secrets_name=config.get('secrets_override'))
        if secrets_name:
            self.template.add_output(Output(service_name + "SecretsName",
                                            Description="AWS secrets manager name to pull the secrets from",
                                            Value=secrets_name))

        task_role = self._add_task_role(service_name, config.get('task_role_attached_managed_policy_arns', []))
        task_execution_role = self._add_task_execution_role(service_name, secrets_name)

        builder = TaskDefinitionBuilder(
            environment=self.env,
            service_name=service_name,
            configuration=config,
            region=self.region,
        )

        td = builder.build_cloudformation_resource(
            container_configurations=container_configurations,
            ecr_image_uri=self.ecr_image_uri,
            fallback_task_role=Ref(task_role),
            fallback_task_execution_role=Ref(task_execution_role),
        )

        self.template.add_resource(td)
        maximum_percent = config['deployment'].get('maximum_percent', 200) if 'deployment' in config else 200
        deployment_configuration = DeploymentConfiguration(MinimumHealthyPercent=100,
                                                           MaximumPercent=int(maximum_percent))
        autoscaling_config = config['autoscaling'] if 'autoscaling' in config else {}
        desired_count = self._get_desired_task_count_for_service(service_name,
                                                                 min_count=int(
                                                                     autoscaling_config.get('min_capacity', 0)))
        request_count_per_target_autoscaling_data = {}
        launch_type_svc = {}

        if 'service_tags' in config:
            launch_type_svc['Tags'] = Tags.from_dict(**config['service_tags'])

        if 'udp_interface' in config:
            lb, target_group_name = self._add_ecs_nlb(service_name, config['udp_interface'], launch_type)
            nlb_enabled = 'nlb_enabled' in config['udp_interface'] and config['udp_interface']['nlb_enabled']

            if nlb_enabled:
                elb, service_listener, nlb_sg = self._add_nlb(service_name, config, target_group_name)
                launch_type_svc['DependsOn'] = service_listener.title
                launch_type_svc['NetworkConfiguration'] = NetworkConfiguration(
                    AwsvpcConfiguration=AwsvpcConfiguration(
                        Subnets=[Ref(self.private_subnet1), Ref(self.private_subnet2)],
                        SecurityGroups=[Ref(nlb_sg)])
                )
                self.template.add_output(
                    Output(
                        service_name + "URL",
                        Description="The URL at which the service is accessible",
                        Value=Sub("udp://${" + elb.name + ".DNSName}")
                    )
                )

            if launch_type == LAUNCH_TYPE_EC2:
                launch_type_svc['PlacementStrategies'] = self.PLACEMENT_STRATEGIES
            svc = Service(
                service_name,
                LoadBalancers=[lb],
                Cluster=self.cluster_name,
                TaskDefinition=Ref(td),
                DesiredCount=desired_count,
                LaunchType=launch_type,
                **launch_type_svc,
            )
            self.template.add_output(
                Output(
                    service_name + 'EcsServiceName',
                    Description='The ECS name which needs to be entered',
                    Value=GetAtt(svc, 'Name')
                )
            )
            self.template.add_resource(svc)
        elif 'http_interface' in config:
            lb, target_group_name, target_group = self._add_ecs_lb(service_name, config, launch_type)

            request_count_per_target_autoscaling_data['target_group'] = target_group

            security_group_ingress = {
                'IpProtocol': 'TCP',
                'ToPort': int(config['http_interface']['container_port']),
                'FromPort': int(config['http_interface']['container_port']),
            }
            create_new_alb = False
            alb = None

            alb_enabled = 'alb' in config['http_interface']
            if alb_enabled:
                alb_config = config['http_interface']['alb']
                create_new_alb = alb_config.get('create_new', False)

                if create_new_alb:
                    alb, service_listener, alb_sg = self._add_alb(service_name, config, target_group_name)
                    launch_type_svc['DependsOn'] = service_listener.title

                    self.template.add_output(
                        Output(
                            service_name + "URL",
                            Description="The URL at which the service is accessible",
                            Value=Sub("https://${" + alb.name + ".DNSName}")
                        )
                    )
                    if launch_type == LAUNCH_TYPE_FARGATE:
                        # needed for FARGATE security group creation.
                        security_group_ingress['SourceSecurityGroupId'] = Ref(alb_sg)
                else:
                    listener_arn = alb_config['listener_arn'] if 'listener_arn' in alb_config \
                        else get_environment_level_alb_listener(self.env)
                    self.attach_to_existing_listener(alb_config, service_name, target_group_name, listener_arn)
                    alb_full_name = self.get_alb_full_name_from_listener_arn(listener_arn)
                    self.create_target_group_alarms(target_group_name, target_group, alb_full_name,
                                                    alb_config)
            if launch_type == LAUNCH_TYPE_FARGATE:
                # if launch type is ec2, then services inherit the ec2 instance security group
                # otherwise, we need to specify a security group for the service
                service_security_group = SecurityGroup(
                    pascalcase("FargateService" + self.env + service_name),
                    GroupName=pascalcase("FargateService" + self.env + service_name),
                    SecurityGroupIngress=[security_group_ingress],
                    VpcId=Ref(self.vpc),
                    GroupDescription=pascalcase("FargateService" + self.env + service_name)
                )
                self.template.add_resource(service_security_group)

                launch_type_svc['NetworkConfiguration'] = NetworkConfiguration(
                    AwsvpcConfiguration=AwsvpcConfiguration(
                        Subnets=[
                            Ref(self.private_subnet1),
                            Ref(self.private_subnet2)
                        ],
                        SecurityGroups=[
                            Ref(service_security_group)
                        ]
                    )
                )
            else:
                launch_type_svc['Role'] = self.ecs_service_role
                launch_type_svc['PlacementStrategies'] = self.PLACEMENT_STRATEGIES

            svc = Service(
                service_name,
                LoadBalancers=[lb],
                Cluster=self.cluster_name,
                TaskDefinition=Ref(td),
                DesiredCount=desired_count,
                DeploymentConfiguration=deployment_configuration,
                LaunchType=launch_type,
                **launch_type_svc,
            )
            if create_new_alb and alb:
                request_count_per_target_autoscaling_data['create_new_alb'] = create_new_alb
                request_count_per_target_autoscaling_data['alb_arn'] = alb

            self.template.add_output(
                Output(
                    service_name + 'EcsServiceName',
                    Description='The ECS name which needs to be entered',
                    Value=GetAtt(svc, 'Name')
                )
            )

            self.template.add_resource(svc)
        elif 'tcp_interface' in config:
            launch_type_svc['NetworkConfiguration'] = NetworkConfiguration(
                AwsvpcConfiguration=AwsvpcConfiguration(
                    Subnets=[Ref(self.private_subnet1), Ref(self.private_subnet2)],
                    SecurityGroups=[config['tcp_interface']['target_security_group']])
            )
            svc = Service(
                service_name,
                LoadBalancers=[
                    LoadBalancer(ContainerName=container_name(service_name),
                                 TargetGroupArn=config['tcp_interface']['target_group_arn'],
                                 ContainerPort=int(config['tcp_interface']['container_port']))],
                Cluster=self.cluster_name,
                TaskDefinition=Ref(td),
                DesiredCount=desired_count,
                LaunchType=launch_type,
                PlacementStrategies=self.PLACEMENT_STRATEGIES,
                **launch_type_svc
            )
            self.template.add_output(
                Output(
                    service_name + 'EcsServiceName',
                    Description='The ECS name which needs to be entered',
                    Value=GetAtt(svc, 'Name')
                )
            )
            self.template.add_resource(svc)
        else:
            if launch_type == LAUNCH_TYPE_FARGATE:
                # if launch type is ec2, then services inherit the ec2 instance security group
                # otherwise, we need to specify a security group for the service
                service_security_group = SecurityGroup(
                    pascalcase("FargateService" + self.env + service_name),
                    GroupName=pascalcase("FargateService" + self.env + service_name),
                    SecurityGroupIngress=[],
                    VpcId=Ref(self.vpc),
                    GroupDescription=pascalcase("FargateService" + self.env + service_name)
                )
                self.template.add_resource(service_security_group)
                launch_type_svc['NetworkConfiguration'] = NetworkConfiguration(
                    AwsvpcConfiguration=AwsvpcConfiguration(
                        Subnets=[
                            Ref(self.private_subnet1),
                            Ref(self.private_subnet2)
                        ],
                        SecurityGroups=[
                            Ref(service_security_group)
                        ]
                    )
                )
            else:
                launch_type_svc['PlacementStrategies'] = self.PLACEMENT_STRATEGIES

            svc = Service(
                service_name,
                Cluster=self.cluster_name,
                TaskDefinition=Ref(td),
                DesiredCount=desired_count,
                DeploymentConfiguration=deployment_configuration,
                LaunchType=launch_type,
                **launch_type_svc
            )
            self.template.add_output(
                Output(
                    service_name + 'EcsServiceName',
                    Description='The ECS name which needs to be entered',
                    Value=GetAtt(svc, 'Name')
                )
            )
            self.template.add_resource(svc)

        if autoscaling_config:
            self._add_autoscaling_resources(service_name, svc, autoscaling_config,
                                            request_count_per_target_autoscaling_data)
        self._add_service_alarms(svc, config)

    def _add_autoscaling_resources(self, service_name, svc, autoscaling_config,
                                   request_count_per_target_autoscaling_data=None):
        scalable_target = self._add_scalable_target(svc, autoscaling_config)

        if autoscaling_config.get('enable_scalable_target_alarms', True):
            self._add_scalable_target_alarms(service_name, svc, autoscaling_config)

        if 'request_count_per_target' in autoscaling_config:
            request_count_per_target_conf = autoscaling_config.get('request_count_per_target')

            if 'alb_arn' in request_count_per_target_conf:
                alb_arn = request_count_per_target_conf['alb_arn']
            elif request_count_per_target_autoscaling_data.get('create_new_alb'):
                alb_arn = request_count_per_target_autoscaling_data['alb_arn']
            else:
                raise UnrecoverableException('Unable to fetch alb arn, please provide alb_arn in config')

            self._add_alb_request_count_scaling_policy(
                svc,
                alb_arn,
                request_count_per_target_autoscaling_data.get('target_group'),
                request_count_per_target_conf,
                scalable_target
            )

        if 'custom_metric' in autoscaling_config:
            self._add_custom_metric_scaling_policy(
                svc,
                autoscaling_config.get('custom_metric'),
                scalable_target
            )

    def get_alb_full_name_from_listener_arn(self, listener_arn):
        return "/".join(listener_arn.split('/')[1:-1])

    def attach_to_existing_listener(self, alb_config, service_name, target_group_name, listener_arn):
        conditions = []
        if 'host' in alb_config:
            conditions.append(
                Condition(
                    Field="host-header",
                    HostHeaderConfig=HostHeaderConfig(
                        Values=[alb_config['host']],
                    ),
                )
            )
        if 'path' in alb_config:
            conditions.append(
                Condition(
                    Field="path-pattern",
                    PathPatternConfig=PathPatternConfig(
                        Values=[alb_config['path']],
                    ),
                )
            )

        priority = alb_config['priority']
        self.template.add_resource(
            ListenerRule(
                service_name + "ListenerRule",
                ListenerArn=listener_arn,
                Priority=int(priority),
                Conditions=conditions,
                Actions=[Action(
                    Type="forward",
                    TargetGroupArn=Ref(target_group_name),
                )]
            )
        )

    def _add_task_role(self, service_name, managed_policy_arns):
        role = Role(
            service_name + "Role",
            ManagedPolicyArns=managed_policy_arns,
            AssumeRolePolicyDocument=PolicyDocument(
                Statement=[
                    Statement(Effect=Allow,
                              Action=[AssumeRole],
                              Principal=Principal("Service", ["ecs-tasks.amazonaws.com"]))
                ]
            )
        )
        self.template.add_resource(role)
        return role

    def _add_task_execution_role(self, service_name, secrets_name):
        automated_injected_secret_name = get_automated_injected_secret_name(self.env,
                                                                            self.application_name,
                                                                            service_name)
        # https://docs.aws.amazon.com/code-samples/latest/catalog/iam_policies-secretsmanager-asm-user-policy-grants-access-to-secret-by-name-with-wildcard.json.html
        allow_secrets = [Statement(Effect=Allow, Action=[GetSecretValue], Resource=[
            f"arn:aws:secretsmanager:{self.region}:{self.account_id}:secret:{automated_injected_secret_name}-??????",
            f"arn:aws:secretsmanager:{self.region}:{self.account_id}:secret:{secrets_name}-??????",
        ])] \
            if secrets_name else []

        task_execution_role = self.template.add_resource(Role(
            service_name + "TaskExecutionRole",
            AssumeRolePolicyDocument=PolicyDocument(
                Statement=[
                    Statement(
                        Effect=Allow,
                        Action=[AssumeRole],
                        Principal=Principal("Service", ["ecs-tasks.amazonaws.com"])
                    )
                ]
            ),
            Policies=[
                Policy(PolicyName=service_name + "TaskExecutionRolePolicy",
                       PolicyDocument=PolicyDocument(
                           Statement=[
                               *allow_secrets,
                               Statement(Effect=Allow,
                                         Action=[GetAuthorizationToken, BatchCheckLayerAvailability,
                                                 GetDownloadUrlForLayer, BatchGetImage, CreateLogStream, PutLogEvents],
                                         Resource=["*"])
                           ]
                       ))]
        ))
        return task_execution_role

    def _add_ecs_lb(self, service_name, config, launch_type):
        target_group_name = "TargetGroup" + service_name
        health_check_path = config['http_interface']['health_check_path'] if 'health_check_path' in config[
            'http_interface'] else "/elb-check"
        if config['http_interface']['internal']:
            target_group_name = target_group_name + 'Internal'

        target_group_config = {}
        if launch_type == LAUNCH_TYPE_FARGATE:
            target_group_config['TargetType'] = 'ip'

        service_target_group = TargetGroup(
            target_group_name,
            HealthCheckPath=health_check_path,
            HealthyThresholdCount=int(config['http_interface'].get('health_check_healthy_threshold_count',
                                                                   DEFAULT_HEALTH_CHECK_HEALTHY_THRESHOLD_COUNT)),
            HealthCheckIntervalSeconds=int(config['http_interface'].get('health_check_interval_seconds',
                                                                        DEFAULT_HEALTH_CHECK_INTERVAL_SECONDS)),
            HealthCheckTimeoutSeconds=int(config['http_interface'].get('health_check_timeout_seconds',
                                                                       DEFAULT_HEALTH_CHECK_TIMEOUT_SECONDS)),
            UnhealthyThresholdCount=int(config['http_interface'].get('health_check_unhealthy_threshold_count',
                                                                     DEFAULT_HEALTH_CHECK_UNHEALTHY_THRESHOLD_COUNT)),
            TargetGroupAttributes=[
                TargetGroupAttribute(
                    Key='deregistration_delay.timeout_seconds',
                    Value=str(config['http_interface'].get('deregistration_delay',
                                                           DEFAULT_TARGET_GROUP_DEREGISTRATION_DELAY))
                ),
                TargetGroupAttribute(
                    Key='load_balancing.algorithm.type',
                    Value=str(
                        config['http_interface'].get('load_balancing_algorithm', DEFAULT_LOAD_BALANCING_ALGORITHM))
                )
            ],
            VpcId=Ref(self.vpc),
            Protocol="HTTP",
            Matcher=Matcher(HttpCode="200-399"),
            Port=int(config['http_interface']['container_port']),
            **target_group_config
        )

        self.template.add_resource(service_target_group)

        lb = LoadBalancer(
            ContainerName=container_name(service_name),
            TargetGroupArn=Ref(service_target_group),
            ContainerPort=int(config['http_interface']['container_port'])
        )

        return lb, target_group_name, service_target_group

    def _add_ecs_nlb(self, service_name, elb_config, launch_type):
        target_group_name = "TargetGroup" + service_name
        health_check_path = elb_config['health_check_path'] if 'health_check_path' in elb_config else "/elb-check"
        if elb_config['internal']:
            target_group_name = target_group_name + 'Internal'

        target_group_config = {'Port': int(elb_config['container_port']),
                               'HealthCheckPort': int(elb_config['health_check_port']), 'TargetType': 'ip'}
        service_target_group = TargetGroup(
            target_group_name,
            Protocol='UDP',
            # Health check healthy threshold and unhealthy
            # threshold must be the same for target groups with the UDP protocol
            HealthyThresholdCount=int(elb_config.get('health_check_healthy_threshold_count',
                                                     DEFAULT_HEALTH_CHECK_HEALTHY_THRESHOLD_COUNT)),
            HealthCheckIntervalSeconds=int(elb_config.get('health_check_interval_seconds',
                                                          DEFAULT_HEALTH_CHECK_INTERVAL_SECONDS)),
            HealthCheckTimeoutSeconds=int(elb_config.get('health_check_timeout_seconds',
                                                         DEFAULT_HEALTH_CHECK_TIMEOUT_SECONDS)),
            UnhealthyThresholdCount=int(elb_config.get('health_check_healthy_threshold_count',
                                                       DEFAULT_HEALTH_CHECK_HEALTHY_THRESHOLD_COUNT)),
            TargetGroupAttributes=[
                TargetGroupAttribute(
                    Key='deregistration_delay.timeout_seconds',
                    Value=str(elb_config.get('deregistration_delay', DEFAULT_TARGET_GROUP_DEREGISTRATION_DELAY))
                )
            ],
            VpcId=Ref(self.vpc),
            **target_group_config
        )

        self.template.add_resource(service_target_group)

        lb = LoadBalancer(
            ContainerName=container_name(service_name),
            TargetGroupArn=Ref(service_target_group),
            ContainerPort=int(elb_config['container_port'])
        )

        return lb, target_group_name

    def _add_alb(self, service_name, config, target_group_name):
        sg_name = 'SG' + self.env + service_name
        svc_alb_sg = SecurityGroup(
            re.sub(r'\W+', '', sg_name),
            GroupName=self.env + '-' + service_name,
            SecurityGroupIngress=self._generate_alb_security_group_ingress(config),
            VpcId=Ref(self.vpc),
            GroupDescription=Sub(service_name + "-alb-sg")
        )
        self.template.add_resource(svc_alb_sg)

        alb_name = service_name + pascalcase(self.env)
        if config['http_interface']['internal']:
            alb_subnets = [
                Ref(self.private_subnet1),
                Ref(self.private_subnet2)
            ]
            scheme = "internal"
            alb_name += 'Internal'
            alb_name = alb_name[:32]
            alb = ALBLoadBalancer(
                'ALB' + service_name,
                Subnets=alb_subnets,
                SecurityGroups=[
                    self.alb_security_group,
                    Ref(svc_alb_sg)
                ],
                Name=alb_name,
                Tags=[
                    {'Value': alb_name, 'Key': 'Name'}
                ],
                Scheme=scheme
            )
        else:
            alb_subnets = [
                Ref(self.public_subnet1),
                Ref(self.public_subnet2)
            ]
            alb_name = alb_name[:32]
            alb = ALBLoadBalancer(
                'ALB' + service_name,
                Subnets=alb_subnets,
                SecurityGroups=[
                    self.alb_security_group,
                    Ref(svc_alb_sg)
                ],
                Name=alb_name,
                Tags=[
                    {'Value': alb_name, 'Key': 'Name'}
                ]
            )

        self.template.add_resource(alb)

        target_group_action = Action(
            TargetGroupArn=Ref(target_group_name),
            Type="forward"
        )
        service_listener = self._add_service_listener(
            service_name,
            target_group_action,
            alb,
            config['http_interface']['internal']
        )
        return alb, service_listener, svc_alb_sg

    def _add_nlb(self, service_name, config, target_group_name):
        sg_name = 'SG' + self.env + service_name
        elb_config = config['udp_interface']
        svc_alb_sg = SecurityGroup(
            re.sub(r'\W+', '', sg_name),
            GroupName=self.env + '-' + service_name,
            SecurityGroupIngress=self._generate_nlb_security_group_ingress(elb_config),
            VpcId=Ref(self.vpc),
            GroupDescription=Sub(service_name + "-alb-sg")
        )
        self.template.add_resource(svc_alb_sg)

        nlb_name = service_name + pascalcase(self.env)
        if elb_config['internal']:
            alb_subnets = [
                Ref(self.private_subnet1),
                Ref(self.private_subnet2)
            ]
            scheme = "internal"
            nlb_name += 'Internal'
        else:
            scheme = 'internet-facing'
            alb_subnets = [
                Ref(self.public_subnet1),
                Ref(self.public_subnet2)
            ]
        subnet_info = {}
        subnet_mappings = []
        if 'eip_allocaltion_id1' in elb_config:
            subnet_mappings.append(
                SubnetMapping(SubnetId=alb_subnets[0], AllocationId=elb_config['eip_allocaltion_id1']))
        else:
            subnet_mappings.append(
                SubnetMapping(SubnetId=alb_subnets[0]))

        if 'eip_allocaltion_id2' in elb_config:
            subnet_mappings.append(
                SubnetMapping(SubnetId=alb_subnets[1], AllocationId=elb_config['eip_allocaltion_id2']))
        else:
            subnet_mappings.append(
                SubnetMapping(SubnetId=alb_subnets[1]))

        subnet_info['SubnetMappings'] = subnet_mappings

        nlb_name = nlb_name[:32]
        nlb = NLBLoadBalancer(
            'NLB' + service_name,
            SecurityGroups=[],
            Name=nlb_name,
            Tags=[
                {'Value': nlb_name, 'Key': 'Name'}
            ],
            Scheme=scheme,
            Type='network',
            **subnet_info
        )

        self.template.add_resource(nlb)

        target_group_action = Action(
            TargetGroupArn=Ref(target_group_name),
            Type="forward"
        )

        service_listener = Listener(
            "LoadBalancerListener" + service_name,
            Protocol="UDP",
            DefaultActions=[target_group_action],
            LoadBalancerArn=Ref(nlb),
            Port=int(config['udp_interface']['container_port']),
        )
        self.template.add_resource(service_listener)

        self._add_nlb_alarms(service_name, nlb)
        return nlb, service_listener, svc_alb_sg

    def _add_service_listener(self, service_name, target_group_action,
                              alb, internal):
        ssl_cert = Certificate(CertificateArn=self.ssl_certificate_arn)
        service_listener = Listener(
            "SslLoadBalancerListener" + service_name,
            Protocol="HTTPS",
            DefaultActions=[target_group_action],
            LoadBalancerArn=Ref(alb),
            Port=443,
            Certificates=[ssl_cert],
            SslPolicy="ELBSecurityPolicy-FS-1-2-Res-2019-08"
        )
        self.template.add_resource(service_listener)
        if internal:
            # Allow HTTP traffic on internal services
            http_service_listener = Listener(
                "LoadBalancerListener" + service_name,
                Protocol="HTTP",
                DefaultActions=[target_group_action],
                LoadBalancerArn=Ref(alb),
                Port=80
            )
            self.template.add_resource(http_service_listener)
        else:
            # Redirect HTTP to HTTPS on external services
            redirection_config = RedirectConfig(
                StatusCode='HTTP_301',
                Protocol='HTTPS',
                Port='443'
            )
            http_redirection_action = Action(
                RedirectConfig=redirection_config,
                Type="redirect"
            )
            http_redirection_listener = Listener(
                "LoadBalancerRedirectionListener" + service_name,
                Protocol="HTTP",
                DefaultActions=[http_redirection_action],
                LoadBalancerArn=Ref(alb),
                Port=80
            )
            self.template.add_resource(http_redirection_listener)
        return service_listener

    def _add_nlb_alarms(self, service_name, nlb):
        unhealthy_alarm = Alarm(
            'NlbUnhealthyHostAlarm' + service_name,
            EvaluationPeriods=1,
            Dimensions=[
                MetricDimension(
                    Name='LoadBalancer',
                    Value=GetAtt(nlb, 'LoadBalancerFullName')
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription='Triggers if any host is marked unhealthy',
            Namespace='AWS/NetworkELB',
            Period=60,
            ComparisonOperator='GreaterThanOrEqualToThreshold',
            Statistic='Sum',
            Threshold='1',
            MetricName='UnHealthyHostCount',
            TreatMissingData='notBreaching'
        )

    def _enrich_alarm_description(self, msg, config):
        playbook_url = config.get('alarm_playbook_url', '')
        if playbook_url == '':
            return msg
        msg = msg.rstrip('. ')
        return f"{msg}. More details can be found from {playbook_url}"

    def create_target_group_alarms(self, target_group_name, target_group, alb_full_name, alb_config):
        unhealthy_alarm = Alarm(
            'TargetGroupUnhealthyHostAlarm' + target_group_name,
            EvaluationPeriods=int(alb_config.get('target_group_unhealthy_host_evaluation_periods', 3)),
            Dimensions=[
                MetricDimension(
                    Name='LoadBalancer',
                    Value=alb_full_name
                ),
                MetricDimension(
                    Name='TargetGroup',
                    Value=GetAtt(target_group, 'TargetGroupFullName')
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription=self._enrich_alarm_description('Triggers if any host is marked unhealthy', alb_config),
            Namespace='AWS/ApplicationELB',
            Period=60,
            ComparisonOperator='GreaterThanOrEqualToThreshold',
            Statistic='Sum',
            Threshold=str(alb_config.get('target_group_unhealthy_host_threshold', '5')),
            MetricName='UnHealthyHostCount',
            TreatMissingData='notBreaching'
        )
        self.template.add_resource(unhealthy_alarm)

        high_5xx_alarm = Alarm(
            'HighTarget5XXAlarm' + target_group_name,
            EvaluationPeriods=int(alb_config.get('target_5xx_error_evaluation_periods', 5)),
            Dimensions=[
                MetricDimension(
                    Name='LoadBalancer',
                    Value=alb_full_name
                ),
                MetricDimension(
                    Name='TargetGroup',
                    Value=GetAtt(target_group, 'TargetGroupFullName')
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription=self._enrich_alarm_description('Triggers if target returns 5xx error code', alb_config),
            Namespace='AWS/ApplicationELB',
            Period=int(alb_config.get('target_5xx_error_period', 60)),
            ComparisonOperator='GreaterThanOrEqualToThreshold',
            Statistic='Sum',
            Threshold=int(alb_config.get('target_5xx_error_threshold', 300)),
            MetricName='HTTPCode_Target_5XX_Count',
            TreatMissingData='notBreaching'
        )
        self.template.add_resource(high_5xx_alarm)

        latency_alarms = self._get_latency_alarms(alb_config, alb_full_name,
                                                  target_group, target_group_name)

        for alarm in latency_alarms:
            self.template.add_resource(alarm)

    def _get_latency_alarms(self, alb_config, alb_full_name, target_group, target_group_name):
        high_p95_latency_alarm = Alarm(
            'HighP95LatencyAlarm' + target_group_name,
            EvaluationPeriods=int(alb_config.get('target_p95_latency_evaluation_periods', 5)),
            Dimensions=[
                MetricDimension(
                    Name='LoadBalancer',
                    Value=alb_full_name
                ),
                MetricDimension(
                    Name='TargetGroup',
                    Value=GetAtt(target_group, 'TargetGroupFullName')
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription=self._enrich_alarm_description('Triggers if p95 latency of target group is higher than threshold', alb_config),
            Namespace='AWS/ApplicationELB',
            Period=int(alb_config.get('target_p95_latency_period_seconds', 60)),
            ComparisonOperator='GreaterThanOrEqualToThreshold',
            ExtendedStatistic='p95',
            Threshold=int(alb_config.get('target_p95_latency_threshold_seconds', 15)),
            MetricName='TargetResponseTime',
            TreatMissingData='notBreaching'
        )
        high_p99_latency_alarm = Alarm(
            'HighP99LatencyAlarm' + target_group_name,
            EvaluationPeriods=int(alb_config.get('target_p99_latency_evaluation_periods', 5)),
            Dimensions=[
                MetricDimension(
                    Name='LoadBalancer',
                    Value=alb_full_name
                ),
                MetricDimension(
                    Name='TargetGroup',
                    Value=GetAtt(target_group, 'TargetGroupFullName')
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription=self._enrich_alarm_description('Triggers if p99 latency of target group is higher than threshold', alb_config),
            Namespace='AWS/ApplicationELB',
            Period=int(alb_config.get('target_p99_latency_period_seconds', 60)),
            ComparisonOperator='GreaterThanOrEqualToThreshold',
            ExtendedStatistic='p99',
            Threshold=int(alb_config.get('target_p99_latency_threshold_seconds', 25)),
            MetricName='TargetResponseTime',
            TreatMissingData='notBreaching'
        )
        return [high_p95_latency_alarm, high_p99_latency_alarm]

    def _add_alb_alarms(self, service_name, alb, alb_config):
        elb_5XX_error_codes_to_monitor = ["502", "503", "504"]
        for elb_5xx_error_code in elb_5XX_error_codes_to_monitor:
            threshold = alb_config.get(f'elb_{elb_5xx_error_code}_error_threshold', "5")
            self.template.add_resource(Alarm(
                f'ELB{elb_5xx_error_code}Count' + service_name,
                EvaluationPeriods=1,
                Dimensions=[
                    MetricDimension(
                        Name='LoadBalancer',
                        Value=GetAtt(alb, 'LoadBalancerFullName')
                    )
                ],
                AlarmActions=[Ref(self.notification_sns_arn)],
                OKActions=[Ref(self.notification_sns_arn)],
                AlarmDescription=self._enrich_alarm_description(f'Triggers if {elb_5xx_error_code} response originated from load balancer', alb_config),
                Namespace='AWS/ApplicationELB',
                Period=60,
                ComparisonOperator='GreaterThanOrEqualToThreshold',
                Statistic='Sum',
                Threshold=threshold,
                MetricName=f'HTTPCode_ELB_{elb_5xx_error_code}_Count',
                TreatMissingData='notBreaching'
            ))

        rejected_connections_alarm = Alarm(
            'ElbRejectedConnectionsAlarm' + service_name,
            EvaluationPeriods=1,
            Dimensions=[
                MetricDimension(
                    Name='LoadBalancer',
                    Value=GetAtt(alb, 'LoadBalancerFullName')
                )
            ],
            AlarmActions=[Ref(self.notification_sns_arn)],
            OKActions=[Ref(self.notification_sns_arn)],
            AlarmDescription=self._enrich_alarm_description('Triggers if load balancer has rejected connections because the load balancer had reached its maximum number of connections.', alb_config),
            Namespace='AWS/ApplicationELB',
            Period=60,
            ComparisonOperator='GreaterThanOrEqualToThreshold',
            Statistic='Sum',
            Threshold='1',
            MetricName='RejectedConnectionCount',
            TreatMissingData='notBreaching'
        )
        self.template.add_resource(rejected_connections_alarm)

    def _generate_alb_security_group_ingress(self, config):
        ingress_rules = []
        for access_ip in config['http_interface']['restrict_access_to']:
            if access_ip.find('/') == -1:
                access_ip = access_ip + '/32'
            ingress_rules.append({
                'ToPort': 80,
                'IpProtocol': 'TCP',
                'FromPort': 80,
                'CidrIp': access_ip
            })
            ingress_rules.append({
                'ToPort': 443,
                'IpProtocol': 'TCP',
                'FromPort': 443,
                'CidrIp': access_ip
            })
        return ingress_rules

    def _generate_nlb_security_group_ingress(self, elb_config):
        ingress_rules = []
        for access_ip in elb_config['restrict_access_to']:
            if access_ip.find('/') == -1:
                access_ip = access_ip + '/32'
            port = elb_config['container_port']
            health_check_port = elb_config['health_check_port']
            ingress_rules.append({
                'ToPort': int(port),
                'IpProtocol': 'UDP',
                'FromPort': int(port),
                'CidrIp': access_ip
            })
            ingress_rules.append(
                {
                    'ToPort': int(health_check_port),
                    'IpProtocol': 'TCP',
                    'FromPort': int(health_check_port),
                    'CidrIp': access_ip
                })
        return ingress_rules

    def _add_ecs_service_iam_role(self):
        injected_service_role_arn = self.configuration.get('service_role_arn')
        role_name = Sub('ecs-svc-${AWS::StackName}-${AWS::Region}')
        assume_role_policy = {
            u'Statement': [
                {
                    u'Action': [u'sts:AssumeRole'],
                    u'Effect': u'Allow',
                    u'Principal': {
                        u'Service': [u'ecs.amazonaws.com']
                    }
                }
            ]
        }
        ecs_service_role = Role(
            'ECSServiceRole',
            Path='/',
            ManagedPolicyArns=[
                'arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceRole'
            ],
            RoleName=role_name,
            AssumeRolePolicyDocument=assume_role_policy
        )
        self.template.add_resource(ecs_service_role)
        if injected_service_role_arn:
            return injected_service_role_arn
        else:
            return Ref(ecs_service_role)

    def _add_service_outputs(self):
        self.template.add_output(Output(
            "CloudliftOptions",
            Description="Options used with cloudlift when \
building this service",
            Value=json.dumps(
                self.configuration,
                cls=DecimalEncoder
            )
        ))
        self._add_stack_outputs()

    def _add_ecr_outputs(self):
        ecr_repo_config = self.configuration.get('ecr_repo')
        self.template.add_output(
            Output(
                "ECRRepoName",
                Description="ECR repo to for docker images",
                Value=ecr_repo_config.get('name')
            )
        )

        if 'account_id' in ecr_repo_config:
            self.template.add_output(
                Output(
                    "ECRAccountID",
                    Description="Account ID to which the ECR repo belongs to",
                    Value=ecr_repo_config.get('account_id')
                )
            )
        if 'assume_role_arn' in ecr_repo_config:
            self.template.add_output(
                Output(
                    "ECRAssumeRoleARN",
                    Description="Role to assume to interact with ECR",
                    Value=ecr_repo_config.get('assume_role_arn')
                )
            )

    def _add_service_parameters(self):
        self.notification_sns_arn = Parameter(
            "NotificationSnsArn",
            Description='',
            Type="String",
            Default=self.notifications_arn)
        self.template.add_parameter(self.notification_sns_arn)
        self.vpc = Parameter(
            "VPC",
            Description='',
            Type="AWS::EC2::VPC::Id",
            Default=list(
                filter(
                    lambda x: x['OutputKey'] == "VPC",
                    self.environment_stack['Outputs']
                )
            )[0]['OutputValue']
        )
        self.template.add_parameter(self.vpc)
        self.public_subnet1 = Parameter(
            "PublicSubnet1",
            Description='',
            Type="AWS::EC2::Subnet::Id",
            Default=list(
                filter(
                    lambda x: x['OutputKey'] == "PublicSubnet1",
                    self.environment_stack['Outputs']
                )
            )[0]['OutputValue']
        )
        self.template.add_parameter(self.public_subnet1)
        self.public_subnet2 = Parameter(
            "PublicSubnet2",
            Description='',
            Type="AWS::EC2::Subnet::Id",
            Default=list(
                filter(
                    lambda x: x['OutputKey'] == "PublicSubnet2",
                    self.environment_stack['Outputs']
                )
            )[0]['OutputValue']
        )
        self.template.add_parameter(self.public_subnet2)
        self.private_subnet1 = Parameter(
            "PrivateSubnet1",
            Description='',
            Type="AWS::EC2::Subnet::Id",
            Default=list(
                filter(
                    lambda x: x['OutputKey'] == "PrivateSubnet1",
                    self.environment_stack['Outputs']
                )
            )[0]['OutputValue']
        )
        self.template.add_parameter(self.private_subnet1)
        self.private_subnet2 = Parameter(
            "PrivateSubnet2",
            Description='',
            Type="AWS::EC2::Subnet::Id",
            Default=list(
                filter(
                    lambda x: x['OutputKey'] == "PrivateSubnet2",
                    self.environment_stack['Outputs']
                )
            )[0]['OutputValue']
        )
        self.template.add_parameter(self.private_subnet2)
        self.template.add_parameter(Parameter(
            "Environment",
            Description='',
            Type="String",
            Default="production"
        ))
        self.alb_security_group = list(
            filter(
                lambda x: x['OutputKey'] == "SecurityGroupAlb",
                self.environment_stack['Outputs']
            )
        )[0]['OutputValue']

    def _get_desired_task_count_for_service(self, service_name, min_count=0):
        return max(self.desired_counts.get(service_name, 1), min_count)

    @property
    def account_id(self):
        return get_account_id()

    @property
    def repo_name(self):
        return self.application_name + '-repo'

    @property
    def notifications_arn(self):
        """
        Get the SNS arn either from service configuration or the cluster
        """
        if 'notifications_arn' in self.configuration:
            return self.configuration['notifications_arn']
        else:
            return TemplateGenerator.notifications_arn.fget(self)


import struct

def to_csv(name, maxdata):
    # (1)  csv 저장할 데이터 준비
    # 레이블 파일과 이미지 파일 열기
    lbl_f = open("../data/MNIST/"+name+"-labels-idx1-ubyte","rb")
    img_f = open("../data/MNIST/"+name+"-images-idx3-ubyte","rb")
    csv_f = open("../data/MNIST/"+name+".csv","w",encoding='utf-8')

    # 헤더 정보읽기
    mag, lbl_count = struct.unpack(">II", lbl_f.read(8)) # 최상위비트부터 unsigned Int형으로
    mag, img_count = struct.unpack(">II", img_f.read(8))
    rows, cols = struct.unpack(">II", img_f.read(8))
    pixels = rows*cols

    # 이미지 데이터를 읽고 csv로 저장하기
    res = []
    for idx in range(lbl_count):
        if idx > maxdata : break

        # 숫자이미지 데이터가 의미하는 숫자값 읽기
        label = struct.unpack("B",lbl_f.read(1))[0] # 리턴타입(튜플)
        # 이미지 데이터 읽기
        bdata = img_f.read(pixels)
        sdata = list(map(lambda n:str(n), bdata))
        csv_f.write(str(label)+",")
        csv_f.write(",".join(sdata)+"\r\n")

        # 잘 저장됐는지 이미지 파일로 저장해서 테스트하기
        if idx<10:
            s = "P2 28 28 255\n"
            s += " ".join(sdata)
            iname = "../data/MNIST/{}-{}-{}.pgm".format(name,idx,label)
            with open(iname,"w",encoding='utf-8') as fp:
                fp.write(s)
    csv_f.close()
    lbl_f.close()
    img_f.close()

# 결과를 파일로 출력
to_csv("train",1000)
to_csv("t10k",500)


import numpy as np
import gym
import pickle

from argparse import ArgumentParser

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from collector import dataCollector
from train import trainer
from test import test

parser = ArgumentParser(description="Parameters for the code - gated trace")
parser.add_argument('--trace_type', type=str, default="gated", help="gated or accumulating" )
parser.add_argument('--gamma', type=float, default=0.99, help="discount factor")
parser.add_argument('--n', type=int, default=8, help="chain length")
parser.add_argument('--t_seeds', type=int, default=25, help="total seeds")
parser.add_argument('--env', type=str, default="gridWorld", help="Environment")
parser.add_argument('--intrst', type=float, default=0.01, help="interest of states while using etd")
parser.add_argument('--hidden', type=int, default=1, help="number of hidden units")
parser.add_argument('--test_every', type=int, default=1, help="calculate MSE after # episode")
parser.add_argument('--episodes', type=int, default=30, help="number of episodes")
parser.add_argument('--lr', type=float, default=0.01, help="learning rate")
parser.add_argument('--log', type=int, default=1, help="compute the results")
parser.add_argument('--save', type=int, default=1, help="save the results")
parser.add_argument('--fo', action='store_true', help="type of env")
args = parser.parse_args()

total_seeds = args.t_seeds

seed_errors = []
for seed in range(total_seeds):
	args.seed = seed
	np.random.seed(args.seed)
	torch.manual_seed(args.seed)
	torch.backends.cudnn.deterministic = True
	torch.backends.cudnn.benchmark = False

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	if args.env == "gridWorld":
		from grid_world import gridWorld
		from networks import gridNet

		env = gridWorld(n=args.n, po = not(args.fo))
		val_net = gridNet(n=args.n, hidden=args.hidden)
		val_net.to(device)

	elif args.env == "gridWorld2":
		from grid_world2 import gridWorld2
		from networks import gridNet

		env = gridWorld2(n=args.n, po = not(args.fo))
		val_net = gridNet(n=args.n, hidden=args.hidden)
		val_net.to(device)

	elif args.env == "lightWorld":
		from light_world import lightWorld
		from networks import lightNet

		env = lightWorld(n=args.n)
		val_net = lightNet(n=args.n)
		val_net.to(device)

	env.seed(args.seed)

	fname_data = "drl_env_"+str(args.env)+"_size_"+str(args.n)+"_seed_"+str(seed)+"_epi_250"
	with open("data/"+fname_data+"_data.pkl", "rb") as f:
		data_list = pickle.load(f)
	optimizer = optim.SGD(val_net.parameters(), lr=args.lr)
	train_cls = trainer(args, data_list, val_net, optimizer, device)
	test_cls = test(args, env, np.ones((env.n**2,env.action_space.n))/env.action_space.n, device)
	val_net, error_list = train_cls.train(test_cls)

	seed_errors.append(error_list)

seed_error_mean = np.array(seed_errors).mean(axis=0)
seed_error_std = np.array(seed_errors).std(axis=0)

if args.log == 1 and args.save == 1:
	if args.trace_type in ["etd", "etd_adaptive"]:
		filename = "drl_"+args.trace_type+"_int_"+str(args.intrst)+"_env_"+str(args.env)+"_size_"+str(args.n)+"_hidden_"+str(args.hidden)+"_lr_"+str(args.lr)+"_seeds_"+str(args.t_seeds)+"_epi_"+str(args.episodes)
	else:
		filename = "drl_"+args.trace_type+"_env_"+str(args.env)+"_size_"+str(args.n)+"_hidden_"+str(args.hidden)+"_lr_"+str(args.lr)+"_seeds_"+str(args.t_seeds)+"_epi_"+str(args.episodes)

	with open("results_"+str(args.env)+"/"+filename+"_errors_mean.pkl", "wb") as f:
		pickle.dump(seed_error_mean, f)

	with open("results_"+str(args.env)+"/"+filename+"_errors_std.pkl", "wb") as f:
		pickle.dump(seed_error_std, f)


from audio2video.__init__ import inp_dir, tar_dir
import argparse
import glob
import os

def init_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--steps',      type=int,   nargs='+',      help='what steps you want to run (0-4)')

    parser.add_argument('--inp_id',     type=str,   default=None,   help='input audio file name')
    parser.add_argument('--tar_id',     type=str,   default=None,   help='target video file name')
    parser.add_argument('--pass_id',    type=str,   default=None,     help='LSTM pass (run) id')

    parser.add_argument('--nthreads',   type=int,   default=10,     help='number of threads when preprocessing')

    parser.add_argument('--train',      type=bool,  default=True,   help='need to train the network or not')
    parser.add_argument('--predict',    type=bool,  default=True,   help='need to predict (infer) sparse mouth shape or not')
    parser.add_argument('--outp_norm',  type=bool,  default=False,  help='normalize sparse mouth shape or not')
    parser.add_argument('--preprocess', type=bool,  default=False,  help='need to re-preprocess data or not')
    parser.add_argument('--vr',         type=float, default=0.2,    help='validation data ratio')
    parser.add_argument('--step_delay', type=int,   default=20,     help='delay step of LSTM')
    parser.add_argument('--dim_hidden', type=int,   default=60,     help='dimension of hidden state')
    parser.add_argument('--nlayers',    type=int,   default=1,      help='number of layers of LSTM')
    parser.add_argument('--keep_prob',  type=float, default=0.8,    help='keep probability')
    parser.add_argument('--seq_len',    type=int,   default=100,    help='length of every training sequence')
    parser.add_argument('--batch_size', type=int,   default=100,    help='training batch size')
    parser.add_argument('--nepochs',    type=int,   default=300,    help='number of epochs')
    parser.add_argument('--grad_clip',  type=int,   default=10,     help='parameter to clip the gradients')
    parser.add_argument('--lr',         type=float, default=1e-3,   help='learning rate')
    parser.add_argument('--dr',         type=float, default=0.99,   help='decay rate of learning rate')
    parser.add_argument('--b_savef',    type=int,   default=50,     help='batch report save frequency')
    parser.add_argument('--e_savef',    type=int,   default=5,      help='checkpoint save frequency')
    parser.add_argument('--argspath',   type=str,   default=None,   help='user-specified args file path')
    parser.add_argument('--showGraph',  type=bool,  default=False,  help='generate tensorboard report or not')

    parser.add_argument('--lineU',      type=int,   default=-1,     help='height of the line seperating upper & lower teeth in upper teeth proxy')
    parser.add_argument('--lineL',      type=int,   default=-1,     help='height of the line seperating upper & lower teeth in lower teeth proxy')

    parser.add_argument('--retiming',   type=bool,  default=True,   help='use the target video after re-timing or the original one')
    return parser

def check_input(opt):
    assert(opt.inp_id is not None)
    flag = os.path.exists('%s/mp3/%s.mp3' % (inp_dir, opt.inp_id))
    assert(flag is True)

def check_target(opt):
    assert(opt.tar_id is not None)
    flag = os.path.exists('%s/mp4/%s.mp4' % (tar_dir, opt.tar_id))
    assert(flag is True)

def check_proxy(opt):
    assert(opt.lineU > 0)
    assert(opt.lineL > 0)
    lower_f = glob.glob('%s/proxy/%s_%s*.png' % (tar_dir, opt.tar_id, 'lower'))
    upper_f = glob.glob('%s/proxy/%s_%s*.png' % (tar_dir, opt.tar_id, 'upper'))
    assert(len(lower_f) == 1 and len(upper_f) == 1)

def parse(parser):
    opt = parser.parse_args()
    opt.steps = sorted(opt.steps)
    assert(min(opt.steps)>=0 and max(opt.steps) <= 4)

    if 0 in opt.steps:
        assert(opt.nthreads > 0)

    if 1 in opt.steps:
        assert(opt.pass_id is not None)
        if opt.predict is True:
            check_input(opt)

    if 2 in opt.steps:
        assert(opt.pass_id is not None)
        check_input(opt)
        check_target(opt)
        check_proxy(opt)

    if 3 in opt.steps:
        check_input(opt)
        check_target(opt)

    if 4 in opt.steps:
        check_input(opt)
        check_target(opt)
    return opt

def ret_opt():
    return parse(init_parser())


"""
File name: database.py
Authos: Ivan Munoz-Gutierrez
Date created: 07/25/2020
Date last modified: 03/13/2021
Python version: 3.9

This file contains functions to be used by fetch_features.py and biosample.py
"""

from Bio import Entrez
from Bio import SeqIO
from datetime import datetime
import sqlite3

def make_list_accessions(infile):
    """
    Reads a txt file that contains a list of UIDs, as accession or BioSample
    numbers, and creates a python list with these UIDs.

    Parameters
    ----------
    infile : string
        Path to txt file that contains UIDs as accession or BioSample numbers.
        Every UID must be separated by a new line character (enter). The list
        can be created in excell by saving the file as txt or using a text
        editor. The txt file must have a header or the function will not
        include the first UID.

    Returns
    -------
    list_accessions: list
        Returns a list of accession numbers. For example, a return list from a
        txt file with nine accession number would look like the following list:
        ['CP049609.1', 'CP028704.1', 'CP043542.1', 'CP040107.1', 'CP041747.1',
         'CP042638.1', 'CP015023.1', 'CP049163.1', 'CP051714.1']
    """
    # Opening infile.txt
    with open(infile, 'r') as reader:
        # Skip the header
        next(reader)
        # Creating a list of accession numbers
        list_accessions = reader.readlines()

    # Remove the '\n' character
    for i, accession_number in enumerate(list_accessions):
        list_accessions[i] = accession_number.replace('\n', '')

    return list_accessions

def mk_uid_batch_list(list_accessions):
    """
    Makes a list of batches of commma-delimited UIDs.

    Every batch created is a set of UIDs separated by commas, and every batch
    has at least the indicaded number of UIDs provided in the batch_size
    variable. Examples of UIDs are accession numbers and BioSample numbers.

    Parameters
    ----------
    list_accessions : list
        List of UIDs to be proccessed

    Returns
    -------
    submission_list : list
        List of strings containg UIDs separated by commas. An example of a
        submission_list created with accession numbers and a batch size of
        three look like the following:
        ['CP049609.1,CP028704.1,CP043542.1',
         'CP040107.1,CP041747.1,CP042638.1',
         'CP015023.1,CP049163.1,CP051714.1']
    """
    # Number of sequences or UIDs to be proccessed
    number_seq = len(list_accessions)
    # Number of sequences or UIDs to be requested by batch. Limit the batch to
    # less than 200.
    batch_size = 5
    # Declaring the list of batches of comma-delimited UIDs that will be return
    # after processing
    submission_list = []
    # Counter to access the list_accessions
    counter_accessions = 0

    # Loop to create the list of UIDs by batches
    for start in range(0, number_seq, batch_size):
        end = min(number_seq, start + batch_size)
        # This list is going to save temporarily the batch of UIDs that are
        # going to be converted into a string of comma-separated UIDs
        set_list = []
        # Making batches
        for _ in range(start, end):
            set_list.append(list_accessions[counter_accessions])
            counter_accessions += 1
        # Converting the list into string
        set_list = ','.join(set_list)
        submission_list.append(set_list)

    return submission_list

def get_biosample_numbers(submission_list, email_address):
    """
    Retrieves BioSample numbers from a list that contains batches of accession
    numbers.

    This function uses Entrez.epost to upload batches of accession numbers to
    the Entrez History server. To avoid problems with large batches of
    accession numbers, limit the number of accession number per batch to 200.

    Parameters
    ----------
    submission_list : list
        List of strings containg batches of accession numbers separated by
        commas. The following is an example of a submission_list containing
        batches of three accession numbers per element in the list:
        ["CP049609.1,CP028704.1,CP043542.1",
         "CP040107.1,CP041747.1,CP042638.1",
         "CP015023.1,CP049163.1,CP051714.1"]
    email_address : string
        User's email address is requested by NCBI

    Returns
    -------
    biosample_numbers : list
        List of BioSample numbers retrieved from Entrez by using the provided
        list of accession numbers. The following is a return list containing
        the corresponding BioSample numbers of the example provided in the
        Parameters section:
        ["SAMN07169263", "SAMN08875353", "SAMEA104140560",
         "SAMN05360217", "SAMN12302771", "SAMN12500846",
         "SAMN04202539", "SAMN14133047", "SAMN14609782"]

    Notes
    -----
    For more information about Entrez.epost read chapter 9.4 of Biopython
    Tutorial and Cookbook (Biopython 1.76) and visit:
    https://www.ncbi.nlm.nih.gov/books/NBK25499/#chapter4.EPost
    """
    # Creating list to save BioSample numbers
    biosample_numbers = []

    # Provide email address to GeneBank
    Entrez.email = email_address

    # Initializing last accession number
    end = 0

    # Fetching BioSample numbers from GenBank by batches
    for _, submission in enumerate(submission_list):
        start = end
        # submission_list is a list of accession numbers separated by
        # commas. Therefore, the number of commas indicate the number of
        # accession numbers.
        batch_size = submission.count(',') + 1
        end = end + batch_size

        # Printing download batch record
        print(f"Retrieving BioSample numbers from record {start + 1} to {end}")

        # Posting the submission_list.
        # Because we are requesting information from a huge list of acc
        # numbers, we have to use the ".epost" function which uploads a
        # list of UIs (acc numbers) for use in subsequent searches.
        # From .epost we can get the QueryKey and the WebEnv which define
        # our history session and can be used to performe searches of data.
        posting = Entrez.epost('nuccore', id=submission)
        search_results = Entrez.read(posting)

        # Copying cookie "WebEnv" and query "QueryKey" from our history
        # session to keep track of our batch fetching.
        # WevEnv -> Web environment string returned from a previous
        # ESearch, EPost or ELink call; QueryKey -> Integer query key
        # returned by a previous ESearch, EPost or ELink call
        webenv = search_results["WebEnv"]
        query_key = search_results["QueryKey"]

        # Getting the batch information
        # db -> database, nuccore -> nuleotide, rettype -> retrieval type
        # retmode -> determines the format of the return output
        # retstart -> sequential index of the first UID in the retrieved
        # set_number to be shown in the XML output
        # retmax -> total number of UIDs from the retrieved set_number to be
        # shown in the XML output
        # idtype-> specifies the type of identifier to return for sequence
        # databases, acc -> accesion number
        fetch_handle = Entrez.efetch(
            db="nuccore",
            rettype="gb",
            retmode="text",
            retstart=0,
            retmax=batch_size,
            webenv=webenv,
            query_key=query_key,
            idtype="acc"
        )

        # Parsing throw the fetched information
        for seq_record in SeqIO.parse(fetch_handle, "gb"):
            # Looping over database cross-references (dbxrefs)
            for xref in enumerate(seq_record.dbxrefs):
                list_dbxrefs = xref[1].split(':')
                # Getting the BioSample number
                if 'BioSample' in list_dbxrefs:
                    biosample_numbers.append(list_dbxrefs[1])
                    break

    return biosample_numbers

def parser(fetch_handle, set_number):
    """
    Function to parse information fetched from nuccore NCBI

    Parameters
    ----------
    fetch_handle : network connection
        Pointer to a network connection that will allow to download and parse
        the requested sequences from the internet
    set_number : inteter
        Number of the batch of sequences that is being analyzed. This number
        will help to keep track of the downloaded sequences.
    seq_counter : integer
        Number of the position of the first accession number in the set. If
        set_number is 1, seq_counter has to be 1.

    Returns
    -------
    records : list
        List of dictionaries. Every dictionary corresponds to fetch features of
        one accession number.
    """
    # Creating a list to save records in dictionaries
    records = []
    info = {}

    # Parsing throw the fetched information
    for seq_record in SeqIO.parse(fetch_handle, "gb"):
        # keeping track of the set_number (or batch) analyzed
        info['set_batch'] = set_number

        # Extracting description
        info['description'] = seq_record.description

        # Extracting sequence id, i.e. accession number
        info['accession'] = seq_record.id

        # '.seq' is an object with the sequence itself
        info['size'] = len(seq_record.seq)

        #######################################################################
        # '.annotations' is a dictionary of aditional information about the
        # sequence as last modification date, topology, sequence_version,
        # organims, references, etc.
        #######################################################################
        # Getting date and give format
        if 'date' in seq_record.annotations:
            mod_date = seq_record.annotations['date']
            mod_date = datetime.strptime(mod_date, '%d-%b-%Y')
            mod_date = datetime.strftime(mod_date, "%Y-%m-%d")
            info['mod_date'] = mod_date
        else:
            info['mod_date'] = 'missing'
        # Getting topology
        if "topology" in seq_record.annotations:
            info['topology'] = seq_record.annotations["topology"]
        else:
            info['topology'] = "missing"

        # Checking whether it is chromosome or plasmid. The 4,000,000 used in
        # size is an arbitrary number that considers chormosomes of that length
        # Getting chromosome
        if 'chromosome' in info['description']:
            info['molecule'] = 'chromosome'
        # Getting chromosome
        elif info['size'] >= 4000000 and info['topology'] == 'circular':
            info['molecule'] = 'chromosome'
        # Getting plasmid
        elif 'plasmid' in info['description']:
            info['molecule'] = 'plasmid'
        else:
            info['molecule'] = 'missing'

        # '.features' is a list of SeqFeatures objects with more structured
        # information about the features on a sequence
        feature = seq_record.features

        # Looping throw list feature
        for index in feature:
            # '.type' is only a description of the type of feature
            # that could be source, CDS, gene, etc.
            # In source we can find organism, strain, host, country, etc.
            if index.type == "source":
                # Creating a dictionary of the qualifiers from source
                dictionary = dict(index.qualifiers)
                # '.get' gives a list
                # Getting molecule type
                if "mol_type" in dictionary:
                    info['mol_type'] = dictionary.get("mol_type")[0]
                else:
                    info['mol_type'] = 'missing'
                # Getting organism
                if 'organism' in dictionary:
                    info['organism'] = dictionary.get('organism')[0]
                else:
                    info['organism'] = 'missing'
                # Getting strain
                if 'strain' in dictionary:
                    info['strain'] = dictionary.get('strain')[0]
                else:
                    info['strain'] = 'missing'
                # Getting isolation source
                if 'isolation_source' in dictionary:
                    info['isolation_source'] =\
                        dictionary.get('isolation_source')[0]
                else:
                    info['isolation_source'] = 'missing'
                # Getting host
                if 'host' in dictionary:
                    info['host'] = dictionary.get('host')[0]
                else:
                    info['host'] = 'missing'
                # Getting plasmid
                if 'plasmid' in dictionary:
                    info['plasmid'] = dictionary.get('plasmid')[0]
                else:
                    info['plasmid'] = 'missing'
                # Getting country
                if 'country' in dictionary:
                    info['country'] = dictionary.get('country')[0]
                else:
                    info['country'] = 'missing'
                # Getting coordinates
                if "lat_lon" in dictionary:
                    info['lat_lon'] = dictionary.get("lat_lon")[0]
                else:
                    info['lat_lon'] = 'missing'
                # Getting collection date
                if "collection_date" in dictionary:
                    info['collection_date'] =\
                        dictionary.get("collection_date")[0]
                else:
                    info['collection_date'] = 'missing'
                # Getting notes
                if "note" in dictionary:
                    info['note'] = dictionary.get("note")[0]
                else:
                    info['note'] = 'missing'
                # Getting serovar
                if "serovar" in dictionary:
                    info['serovar'] = dictionary.get("serovar")[0]
                else:
                    info['serovar'] = 'missing'
                # Getting collected by
                if "collected_by" in dictionary:
                    info['collected_by'] = dictionary.get("collected_by")[0]
                else:
                    info['collected_by'] = 'missing'
                # Getting genotype
                if "genotype" in dictionary:
                    info['genotype'] = dictionary.get("genotype")[0]
                else:
                    info['genotype'] = 'missing'

                break

        # '.dbxrefs' is a list populated from any PROJECT or DBLINK
        # Checking if .dbxrefs has any information
        if len(seq_record.dbxrefs) == 0:
            info['BioProject'] = 'missing'
            info['BioSample'] = 'missing'

        # Converting the list dbxrefs into dictionary
        dictionary_dbxrefs = {}
        for i in range(len(seq_record.dbxrefs)):
            s = seq_record.dbxrefs[i].split(":")
            dictionary_dbxrefs[s[0]] = s[1]

        # Saving BioProject and BioSample
        if "BioProject" in dictionary_dbxrefs:
            info['BioProject'] = dictionary_dbxrefs.get("BioProject")
        else:
            info['BioProject'] = 'missing'
        if "BioSample" in dictionary_dbxrefs:
            info['BioSample'] = dictionary_dbxrefs.get("BioSample")
        else:
            info['BioSample'] = 'missing'

        # Getting the Genome-Assembly-Data
        # Checking if the sequence has structured_comment
        if (("structured_comment" in seq_record.annotations)
            and ("Genome-Assembly-Data" in (
                seq_record.annotations["structured_comment"]))):
            # Getting Assembly Method
            if ("Assembly Method" in (
                    seq_record.annotations["structured_comment"]
                                          ["Genome-Assembly-Data"])):
                info['Assem_Method'] = (
                    seq_record.annotations["structured_comment"]
                                          ["Genome-Assembly-Data"]
                                          ["Assembly Method"])
            else:
                info['Assem_Method'] = "missing"
            # Getting Genome Coverage
            if "Genome Coverage" in (
                seq_record.annotations["structured_comment"]
                                      ["Genome-Assembly-Data"]):
                info['Gen_Coverage'] = (
                    seq_record.annotations["structured_comment"]
                                          ["Genome-Assembly-Data"]
                                          ["Genome Coverage"])
            else:
                info['Gen_Coverage'] = "missing"
            # Getting Sequencing Technology
            if "Sequencing Technology" in (
                seq_record.annotations["structured_comment"]
                                      ["Genome-Assembly-Data"]):
                info['Seq_Technol'] = (
                    seq_record.annotations["structured_comment"]
                                          ["Genome-Assembly-Data"]
                                          ["Sequencing Technology"])
            else:
                info['Seq_Technol'] = "missing"
        # Extracting data from structured comment
        elif "structured_comment" in seq_record.annotations and (
            "Assembly-Data" in (
                seq_record.annotations["structured_comment"])):
            # Getting Assembly Method
            if "Assembly Method" in (
                seq_record.annotations["structured_comment"]
                                      ["Assembly-Data"]):
                info['Assem_Method'] = (
                    seq_record.annotations["structured_comment"]
                                          ["Assembly-Data"]
                                          ["Assembly Method"])
            else:
                info['Assem_Method'] = "missing"
            # Getting Genome Coverage
            if "Genome Coverage" in (
                seq_record.annotations["structured_comment"]
                                      ["Assembly-Data"]):
                info['Gen_Coverage'] = (
                    seq_record.annotations["structured_comment"]
                                          ["Assembly-Data"]
                                          ["Genome Coverage"])
            else:
                info['Gen_Coverage'] = "missing"
            # Getting Sequencing Technology
            if "Sequencing Technology" in (
                seq_record.annotations["structured_comment"]
                                      ["Assembly-Data"]):
                info['Seq_Technol'] = (
                    seq_record.annotations["structured_comment"]
                                          ["Assembly-Data"]
                                          ["Sequencing Technology"])
            else:
                info['Seq_Technol'] = "missing"

        else:
            info['Assem_Method'] = "missing"
            info['Gen_Coverage'] = "missing"
            info['Seq_Technol'] = "missing"

        # Copying info (dictionary) into records (list)
        records.append(info.copy())

    return records

def clean_features(raw_data):
    """
    Cleans up the features obtained from a BioSample number and get the most
    updated information

    Parameters
    ----------
    raw_data : list
        List of dictionaries holding all the data downloaded from a specific
        BioSample associated accession numbers.

    Returns
    -------
    recods : list
        List of dictionaries holding the most updated information from
        BioSample associated accession numbers.
    """

    # Creating empty features.db
    open('features.db', 'w').close()

    # Opening features.db for sqlite3
    conn = sqlite3.connect('features.db')

    # Creating a cursor
    c = conn.cursor()

    # Creating table in SQL for the features of the raw data fetched from NCBI
    c.execute("""CREATE TABLE features_raw (
                set_batch int,
                description text,
                accession text,
                size integer,
                molecule text,
                mod_date text,
                topology text,
                mol_type text,
                organism text,
                strain text,
                isolation_source text,
                host text,
                plasmid text,
                country text,
                lat_lon text,
                collection_date text,
                note text,
                serovar text,
                collected_by text,
                genotype text,
                BioProject text,
                BioSample text,
                Assem_Method text,
                Gen_Coverage text,
                Seq_Technol text
                )""")
    conn.commit()

    # Transfering results obtained from NCBI into the table features_raw
    # created in SQL
    for _, raw_result in enumerate(raw_data):
        c.execute("""INSERT INTO features_raw
                     VALUES(:set_batch, :description, :accession,
                            :size, :molecule, :mod_date, :topology, :mol_type,
                            :organism, :strain, :isolation_source, :host,
                            :plasmid, :country, :lat_lon, :collection_date,
                            :note, :serovar, :collected_by, :genotype,
                            :BioProject, :BioSample, :Assem_Method,
                            :Gen_Coverage, :Seq_Technol)""",
                  {'set_batch': int(raw_result['set_batch']),
                   'description': raw_result['description'],
                   'accession': raw_result['accession'],
                   'size': int(raw_result['size']),
                   'molecule': raw_result['molecule'],
                   'mod_date': raw_result['mod_date'],
                   'topology': raw_result['topology'],
                   'mol_type': raw_result['mol_type'],
                   'organism': raw_result['organism'],
                   'strain': raw_result['strain'],
                   'isolation_source': raw_result['isolation_source'],
                   'host': raw_result['host'],
                   'plasmid': raw_result['plasmid'],
                   'country': raw_result['country'],
                   'lat_lon': raw_result['lat_lon'],
                   'collection_date': raw_result['collection_date'],
                   'note': raw_result['note'],
                   'serovar': raw_result['serovar'],
                   'collected_by': raw_result['collected_by'],
                   'genotype': raw_result['genotype'],
                   'BioProject': raw_result['BioProject'],
                   'BioSample': raw_result['BioSample'],
                   'Assem_Method': raw_result['Assem_Method'],
                   'Gen_Coverage': raw_result['Gen_Coverage'],
                   'Seq_Technol': raw_result['Seq_Technol']})
    conn.commit()

    # Getting the most updated files and ordering by size
    c.execute(
        """
        SELECT *
          FROM (
               SELECT table_one.*
                 FROM features_raw table_one
                WHERE table_one.mod_date = (
                      SELECT MAX(table_two.mod_date)
                        FROM features_raw table_two
                       WHERE table_two.size = table_one.size))
         WHERE molecule != "missing"
         ORDER BY size DESC
        """)

    # Extracting the updated results from SQL. The fetchall will create a list
    # of tuples
    results = c.fetchall()

    # Converting results (list of tuples) into a list of dictionaries
    records = []
    info = {}
    for _, updated_result in enumerate(results):
        info['set_batch'] = int(updated_result[0])
        info['description'] = updated_result[1]
        info['accession'] = updated_result[2]
        info['size'] = int(updated_result[3])
        info['molecule'] = updated_result[4]
        info['mod_date'] = updated_result[5]
        info['topology'] = updated_result[6]
        info['mol_type'] = updated_result[7]
        info['organism'] = updated_result[8]
        info['strain'] = updated_result[9]
        info['isolation_source'] = updated_result[10]
        info['host'] = updated_result[11]
        info['plasmid'] = updated_result[12]
        info['country'] = updated_result[13]
        info['lat_lon'] = updated_result[14]
        info['collection_date'] = updated_result[15]
        info['note'] = updated_result[16]
        info['serovar'] = updated_result[17]
        info['collected_by'] = updated_result[18]
        info['genotype'] = updated_result[19]
        info['BioProject'] = updated_result[20]
        info['BioSample'] = updated_result[21]
        info['Assem_Method'] = updated_result[22]
        info['Gen_Coverage'] = updated_result[23]
        info['Seq_Technol'] = updated_result[24]
        records.append(info.copy())
    conn.close()

    return records


class Solution:
    # @param {integer[]} nums
    # @return {boolean}

    #initial approach. Exceed time limit  O(n^2)
    def containsDuplicate1(self, nums):
        l = []

        for i in nums:
            if i not in l:
                l.append(i)
            else:
                return True

        return False

    #second approach. Sorted and compare the current and the next. O(n)
    def containsDuplicate(self, nums):
        nums.sort()
        for i in range(len(nums)-1):
            if nums[i] == nums[i+1]:
                return True

            return False