from typing import ClassVar, Dict, Optional, List, Type

from attrs import define, field

from resoto_plugin_aws.resource.base import AwsResource, AwsApiSpec, GraphBuilder
from resoto_plugin_aws.resource.ec2 import AwsEc2Instance
from resoto_plugin_aws.utils import ToDict
from resotolib.baseresources import BaseAutoScalingGroup, ModelReference
from resotolib.json_bender import Bender, S, Bend, ForallBend
from resotolib.types import Json
from resoto_plugin_aws.aws_client import AwsClient

@define(eq=False, slots=False)
class AwsAutoScalingLaunchTemplateSpecification:
    kind: ClassVar[str] = "aws_autoscaling_launch_template_specification"
    mapping: ClassVar[Dict[str, Bender]] = {
        "launch_template_id": S("LaunchTemplateId"),
        "launch_template_name": S("LaunchTemplateName"),
        "version": S("Version"),
    }
    launch_template_id: Optional[str] = field(default=None)
    launch_template_name: Optional[str] = field(default=None)
    version: Optional[str] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingMinMax:
    kind: ClassVar[str] = "aws_autoscaling_min_max"
    mapping: ClassVar[Dict[str, Bender]] = {"min": S("Min"), "max": S("Max")}
    min: Optional[int] = field(default=None)
    max: Optional[int] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingInstanceRequirements:
    kind: ClassVar[str] = "aws_autoscaling_instance_requirements"
    mapping: ClassVar[Dict[str, Bender]] = {
        "v_cpu_count": S("VCpuCount") >> Bend(AwsAutoScalingMinMax.mapping),
        "memory_mi_b": S("MemoryMiB") >> Bend(AwsAutoScalingMinMax.mapping),
        "cpu_manufacturers": S("CpuManufacturers", default=[]),
        "memory_gi_b_per_v_cpu": S("MemoryGiBPerVCpu") >> Bend(AwsAutoScalingMinMax.mapping),
        "excluded_instance_types": S("ExcludedInstanceTypes", default=[]),
        "instance_generations": S("InstanceGenerations", default=[]),
        "spot_max_price_percentage_over_lowest_price": S("SpotMaxPricePercentageOverLowestPrice"),
        "on_demand_max_price_percentage_over_lowest_price": S("OnDemandMaxPricePercentageOverLowestPrice"),
        "bare_metal": S("BareMetal"),
        "burstable_performance": S("BurstablePerformance"),
        "require_hibernate_support": S("RequireHibernateSupport"),
        "network_interface_count": S("NetworkInterfaceCount") >> Bend(AwsAutoScalingMinMax.mapping),
        "local_storage": S("LocalStorage"),
        "local_storage_types": S("LocalStorageTypes", default=[]),
        "total_local_storage_gb": S("TotalLocalStorageGB") >> Bend(AwsAutoScalingMinMax.mapping),
        "baseline_ebs_bandwidth_mbps": S("BaselineEbsBandwidthMbps") >> Bend(AwsAutoScalingMinMax.mapping),
        "accelerator_types": S("AcceleratorTypes", default=[]),
        "accelerator_count": S("AcceleratorCount") >> Bend(AwsAutoScalingMinMax.mapping),
        "accelerator_manufacturers": S("AcceleratorManufacturers", default=[]),
        "accelerator_names": S("AcceleratorNames", default=[]),
        "accelerator_total_memory_mi_b": S("AcceleratorTotalMemoryMiB") >> Bend(AwsAutoScalingMinMax.mapping),
    }
    v_cpu_count: Optional[AwsAutoScalingMinMax] = field(default=None)
    memory_mi_b: Optional[AwsAutoScalingMinMax] = field(default=None)
    cpu_manufacturers: List[str] = field(factory=list)
    memory_gi_b_per_v_cpu: Optional[AwsAutoScalingMinMax] = field(default=None)
    excluded_instance_types: List[str] = field(factory=list)
    instance_generations: List[str] = field(factory=list)
    spot_max_price_percentage_over_lowest_price: Optional[int] = field(default=None)
    on_demand_max_price_percentage_over_lowest_price: Optional[int] = field(default=None)
    bare_metal: Optional[str] = field(default=None)
    burstable_performance: Optional[str] = field(default=None)
    require_hibernate_support: Optional[bool] = field(default=None)
    network_interface_count: Optional[AwsAutoScalingMinMax] = field(default=None)
    local_storage: Optional[str] = field(default=None)
    local_storage_types: List[str] = field(factory=list)
    total_local_storage_gb: Optional[AwsAutoScalingMinMax] = field(default=None)
    baseline_ebs_bandwidth_mbps: Optional[AwsAutoScalingMinMax] = field(default=None)
    accelerator_types: List[str] = field(factory=list)
    accelerator_count: Optional[AwsAutoScalingMinMax] = field(default=None)
    accelerator_manufacturers: List[str] = field(factory=list)
    accelerator_names: List[str] = field(factory=list)
    accelerator_total_memory_mi_b: Optional[AwsAutoScalingMinMax] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingLaunchTemplateOverrides:
    kind: ClassVar[str] = "aws_autoscaling_launch_template_overrides"
    mapping: ClassVar[Dict[str, Bender]] = {
        "instance_type": S("InstanceType"),
        "weighted_capacity": S("WeightedCapacity"),
        "launch_template_specification": S("LaunchTemplateSpecification")
        >> Bend(AwsAutoScalingLaunchTemplateSpecification.mapping),
        "instance_requirements": S("InstanceRequirements") >> Bend(AwsAutoScalingInstanceRequirements.mapping),
    }
    instance_type: Optional[str] = field(default=None)
    weighted_capacity: Optional[str] = field(default=None)
    launch_template_specification: Optional[AwsAutoScalingLaunchTemplateSpecification] = field(default=None)
    instance_requirements: Optional[AwsAutoScalingInstanceRequirements] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingLaunchTemplate:
    kind: ClassVar[str] = "aws_autoscaling_launch_template"
    mapping: ClassVar[Dict[str, Bender]] = {
        "launch_template_specification": S("LaunchTemplateSpecification")
        >> Bend(AwsAutoScalingLaunchTemplateSpecification.mapping),
        "overrides": S("Overrides", default=[]) >> ForallBend(AwsAutoScalingLaunchTemplateOverrides.mapping),
    }
    launch_template_specification: Optional[AwsAutoScalingLaunchTemplateSpecification] = field(default=None)
    overrides: List[AwsAutoScalingLaunchTemplateOverrides] = field(factory=list)

@define(eq=False, slots=False)
class AwsAutoScalingInstancesDistribution:
    kind: ClassVar[str] = "aws_autoscaling_instances_distribution"
    mapping: ClassVar[Dict[str, Bender]] = {
        "on_demand_allocation_strategy": S("OnDemandAllocationStrategy"),
        "on_demand_base_capacity": S("OnDemandBaseCapacity"),
        "on_demand_percentage_above_base_capacity": S("OnDemandPercentageAboveBaseCapacity"),
        "spot_allocation_strategy": S("SpotAllocationStrategy"),
        "spot_instance_pools": S("SpotInstancePools"),
        "spot_max_price": S("SpotMaxPrice"),
    }
    on_demand_allocation_strategy: Optional[str] = field(default=None)
    on_demand_base_capacity: Optional[int] = field(default=None)
    on_demand_percentage_above_base_capacity: Optional[int] = field(default=None)
    spot_allocation_strategy: Optional[str] = field(default=None)
    spot_instance_pools: Optional[int] = field(default=None)
    spot_max_price: Optional[str] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingMixedInstancesPolicy:
    kind: ClassVar[str] = "aws_autoscaling_mixed_instances_policy"
    mapping: ClassVar[Dict[str, Bender]] = {
        "launch_template": S("LaunchTemplate") >> Bend(AwsAutoScalingLaunchTemplate.mapping),
        "instances_distribution": S("InstancesDistribution") >> Bend(AwsAutoScalingInstancesDistribution.mapping),
    }
    launch_template: Optional[AwsAutoScalingLaunchTemplate] = field(default=None)
    instances_distribution: Optional[AwsAutoScalingInstancesDistribution] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingInstance:
    kind: ClassVar[str] = "aws_autoscaling_instance"
    mapping: ClassVar[Dict[str, Bender]] = {
        "instance_id": S("InstanceId"),
        "instance_type": S("InstanceType"),
        "availability_zone": S("AvailabilityZone"),
        "lifecycle_state": S("LifecycleState"),
        "health_status": S("HealthStatus"),
        "launch_configuration_name": S("LaunchConfigurationName"),
        "launch_template": S("LaunchTemplate") >> Bend(AwsAutoScalingLaunchTemplateSpecification.mapping),
        "protected_from_scale_in": S("ProtectedFromScaleIn"),
        "weighted_capacity": S("WeightedCapacity"),
    }
    instance_id: Optional[str] = field(default=None)
    instance_type: Optional[str] = field(default=None)
    availability_zone: Optional[str] = field(default=None)
    lifecycle_state: Optional[str] = field(default=None)
    health_status: Optional[str] = field(default=None)
    launch_configuration_name: Optional[str] = field(default=None)
    launch_template: Optional[AwsAutoScalingLaunchTemplateSpecification] = field(default=None)
    protected_from_scale_in: Optional[bool] = field(default=None)
    weighted_capacity: Optional[str] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingSuspendedProcess:
    kind: ClassVar[str] = "aws_autoscaling_suspended_process"
    mapping: ClassVar[Dict[str, Bender]] = {
        "process_name": S("ProcessName"),
        "suspension_reason": S("SuspensionReason"),
    }
    process_name: Optional[str] = field(default=None)
    suspension_reason: Optional[str] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingEnabledMetric:
    kind: ClassVar[str] = "aws_autoscaling_enabled_metric"
    mapping: ClassVar[Dict[str, Bender]] = {"metric": S("Metric"), "granularity": S("Granularity")}
    metric: Optional[str] = field(default=None)
    granularity: Optional[str] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingWarmPoolConfiguration:
    kind: ClassVar[str] = "aws_autoscaling_warm_pool_configuration"
    mapping: ClassVar[Dict[str, Bender]] = {
        "max_group_prepared_capacity": S("MaxGroupPreparedCapacity"),
        "min_size": S("MinSize"),
        "pool_state": S("PoolState"),
        "status": S("Status"),
        "instance_reuse_policy": S("InstanceReusePolicy", "ReuseOnScaleIn"),
    }
    max_group_prepared_capacity: Optional[int] = field(default=None)
    min_size: Optional[int] = field(default=None)
    pool_state: Optional[str] = field(default=None)
    status: Optional[str] = field(default=None)
    instance_reuse_policy: Optional[bool] = field(default=None)

@define(eq=False, slots=False)
class AwsAutoScalingGroup(AwsResource, BaseAutoScalingGroup):
    kind: ClassVar[str] = "aws_autoscaling_group"
    api_spec: ClassVar[AwsApiSpec] = AwsApiSpec("autoscaling", "describe-auto-scaling-groups", "AutoScalingGroups")
    reference_kinds: ClassVar[ModelReference] = {
        "successors": {"default": ["aws_ec2_instance"]},
        "predecessors": {"delete": ["aws_ec2_instance"]},
    }
    mapping: ClassVar[Dict[str, Bender]] = {
        "id": S("AutoScalingGroupName"),
        "tags": S("Tags", default=[]) >> ToDict(),
        "name": S("AutoScalingGroupName"),
        "ctime": S("CreatedTime"),
        "arn": S("AutoScalingGroupARN"),
        "autoscaling_launch_configuration_name": S("LaunchConfigurationName"),
        "autoscaling_launch_template": S("LaunchTemplate") >> Bend(AwsAutoScalingLaunchTemplateSpecification.mapping),
        "autoscaling_mixed_instances_policy": S("MixedInstancesPolicy")
        >> Bend(AwsAutoScalingMixedInstancesPolicy.mapping),
        "min_size": S("MinSize"),
        "max_size": S("MaxSize"),
        "autoscaling_desired_capacity": S("DesiredCapacity"),
        "autoscaling_predicted_capacity": S("PredictedCapacity"),
        "autoscaling_default_cooldown": S("DefaultCooldown"),
        "autoscaling_availability_zones": S("AvailabilityZones", default=[]),
        "autoscaling_load_balancer_names": S("LoadBalancerNames", default=[]),
        "autoscaling_target_group_ar_ns": S("TargetGroupARNs", default=[]),
        "autoscaling_health_check_type": S("HealthCheckType"),
        "autoscaling_health_check_grace_period": S("HealthCheckGracePeriod"),
        "autoscaling_instances": S("Instances", default=[]) >> ForallBend(AwsAutoScalingInstance.mapping),
        "autoscaling_suspended_processes": S("SuspendedProcesses", default=[])
        >> ForallBend(AwsAutoScalingSuspendedProcess.mapping),
        "autoscaling_placement_group": S("PlacementGroup"),
        "autoscaling_vpc_zone_identifier": S("VPCZoneIdentifier"),
        "autoscaling_enabled_metrics": S("EnabledMetrics", default=[])
        >> ForallBend(AwsAutoScalingEnabledMetric.mapping),
        "autoscaling_status": S("Status"),
        "autoscaling_termination_policies": S("TerminationPolicies", default=[]),
        "autoscaling_new_instances_protected_from_scale_in": S("NewInstancesProtectedFromScaleIn"),
        "autoscaling_service_linked_role_arn": S("ServiceLinkedRoleARN"),
        "autoscaling_max_instance_lifetime": S("MaxInstanceLifetime"),
        "autoscaling_capacity_rebalance": S("CapacityRebalance"),
        "autoscaling_warm_pool_configuration": S("WarmPoolConfiguration")
        >> Bend(AwsAutoScalingWarmPoolConfiguration.mapping),
        "autoscaling_warm_pool_size": S("WarmPoolSize"),
        "autoscaling_context": S("Context"),
        "autoscaling_desired_capacity_type": S("DesiredCapacityType"),
        "autoscaling_default_instance_warmup": S("DefaultInstanceWarmup"),
    }
    autoscaling_launch_configuration_name: Optional[str] = field(default=None)
    autoscaling_launch_template: Optional[AwsAutoScalingLaunchTemplateSpecification] = field(default=None)
    autoscaling_mixed_instances_policy: Optional[AwsAutoScalingMixedInstancesPolicy] = field(default=None)
    autoscaling_predicted_capacity: Optional[int] = field(default=None)
    autoscaling_default_cooldown: Optional[int] = field(default=None)
    autoscaling_availability_zones: List[str] = field(factory=list)
    autoscaling_load_balancer_names: List[str] = field(factory=list)
    autoscaling_target_group_ar_ns: List[str] = field(factory=list)
    autoscaling_health_check_type: Optional[str] = field(default=None)
    autoscaling_health_check_grace_period: Optional[int] = field(default=None)
    autoscaling_instances: List[AwsAutoScalingInstance] = field(factory=list)
    autoscaling_suspended_processes: List[AwsAutoScalingSuspendedProcess] = field(factory=list)
    autoscaling_placement_group: Optional[str] = field(default=None)
    autoscaling_vpc_zone_identifier: Optional[str] = field(default=None)
    autoscaling_enabled_metrics: List[AwsAutoScalingEnabledMetric] = field(factory=list)
    autoscaling_status: Optional[str] = field(default=None)
    autoscaling_termination_policies: List[str] = field(factory=list)
    autoscaling_new_instances_protected_from_scale_in: Optional[bool] = field(default=None)
    autoscaling_service_linked_role_arn: Optional[str] = field(default=None)
    autoscaling_max_instance_lifetime: Optional[int] = field(default=None)
    autoscaling_capacity_rebalance: Optional[bool] = field(default=None)
    autoscaling_warm_pool_configuration: Optional[AwsAutoScalingWarmPoolConfiguration] = field(default=None)
    autoscaling_warm_pool_size: Optional[int] = field(default=None)
    autoscaling_context: Optional[str] = field(default=None)
    autoscaling_desired_capacity_type: Optional[str] = field(default=None)
    autoscaling_default_instance_warmup: Optional[int] = field(default=None)

    def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
        for instance in self.autoscaling_instances:
            builder.dependant_node(self, clazz=AwsEc2Instance, id=instance.instance_id)

    def update_resource_tag(self, client: AwsClient, key: str, value: str) -> bool:
        client.call(
            aws_service="autoscaling",
            action="create-or-update-tags",
            result_name=None,
            Tags=[
                {
                    "ResourceId": self.name,
                    "ResourceType": "auto-scaling-group",
                    "Key": key,
                    "Value": value,
                    "PropagateAtLaunch": False,
                }
            ],
        )
        return True

    def delete_resource_tag(self, client: AwsClient, key: str) -> bool:
        client.call(
            aws_service="autoscaling",
            action="delete-tags",
            result_name=None,
            Tags=[
                {
                    "ResourceId": self.name,
                    "ResourceType": "auto-scaling-group",
                    "Key": key,
                }
            ],
        )
        return True

    def delete_resource(self, client: AwsClient) -> bool:
        client.call(
            aws_service=self.api_spec.service,
            action="delete-auto-scaling-group",
            result_name=None,
            AutoScalingGroupName=self.name,
            ForceDelete=True,
        )
        return True

    @classmethod
    def called_mutator_apis(cls) -> List[AwsApiSpec]:
        return [
            AwsApiSpec("autoscaling", "create-or-update-tags"),
            AwsApiSpec("autoscaling", "delete-tags"),
            AwsApiSpec("autoscaling", "delete-auto-scaling-group"),
        ]

resources: List[Type[AwsResource]] = [AwsAutoScalingGroup]


import ee

def mask_landsat_sr(masks, T_threshold=None, omission=True):

    dict_mask = {'cloud': ee.Number(2).pow(5).int(),
                 'cshadow': ee.Number(2).pow(3).int(),
                 'snow': ee.Number(2).pow(4).int()}

    if T_threshold is not None:

        if omission:
            sel_masks = [dict_mask[x] for x in masks if x != 'cloud']
            bits = ee.Number(1)
            for m in sel_masks:
                bits = ee.Number(bits.add(m))

            def wrap(img):
                qa = img.select('pixel_qa')

                cloudbit = ee.Number(2).pow(5).int()
                cloudmask = (qa.bitwiseAnd(cloudbit).eq(0)).Not()

                bt = img.select('TIR')
                temp_th = (T_threshold + 273.15) * 10
                bt_mask = bt.lt(temp_th)  # 1 cloud, 0 clear

                bt_mask = (cloudmask.multiply(bt_mask)).Not()
                mask = (qa.bitwiseAnd(bits).eq(0)).And(bt_mask)

                return img.updateMask(mask) \
                    .copyProperties(source=img).set('system:time_start', img.get('system:time_start'))
        # comission
        else:
            sel_masks = [dict_mask[x] for x in masks]
            bits = ee.Number(1)
            for m in sel_masks:
                bits = ee.Number(bits.add(m))

            def wrap(img):
                qa = img.select('pixel_qa')
                mask = qa.bitwiseAnd(bits).eq(0)

                bt = img.select('TIR')
                temp_th = (T_threshold + 273.15) * 10
                bt_mask = bt.gt(temp_th)  # 1 clear, 0 cloud

                mask = mask.multiply(bt_mask)

                return img.updateMask(mask) \
                    .copyProperties(source=img).set('system:time_start', img.get('system:time_start'))
    else:
        sel_masks = [dict_mask[x] for x in masks]
        bits = ee.Number(1)
        for m in sel_masks:
            bits = ee.Number(bits.add(m))
        def wrap(img):
            qa = img.select('pixel_qa')
            mask = qa.bitwiseAnd(bits).eq(0)
            return img.updateMask(mask) \
                      .copyProperties(source=img).set('system:time_start', img.get('system:time_start'))
    return wrap

def mask_s2_scl(img):
    #mask_kernel = ee.Kernel.square(radius=5)
    scl = img.select('SCL')
    mask = scl.neq(3).And(scl.neq(7)).And(
            scl.neq(8)).And(
            scl.neq(9)).And(
            scl.neq(10)).And(
            scl.neq(11))
    #mask = mask.focal_mode(kernel=mask_kernel, iterations=1).rename('CLOUD')
    return img.addBands(mask).updateMask(mask)\
              .copyProperties(source=img).set('system:time_start', img.get('system:time_start'))

def mask_s2_cdi(cdi=-0.5):
    def wrap(img):
        img_cdi = ee.Algorithms.Sentinel2.CDI(img)
        mask = img_cdi.gt(cdi).rename("mask")
        return img.addBands(mask).updateMask(mask) \
           .copyProperties(source=img).set('system:time_start', img.get('system:time_start'))
    return wrap

def mask_s2(img):
    qa = img.select('QA60')
    cloudBitMask = ee.Number(2).pow(10).int()
    cirrusBitMask = ee.Number(2).pow(11).int()
    mask = qa.bitwiseAnd(cloudBitMask).eq(0).And(qa.bitwiseAnd(cirrusBitMask).eq(0)).rename('CLOUD')
    return img.addBands(mask).updateMask(mask)\
              .copyProperties(source=img).set('system:time_start', img.get('system:time_start'))

def focal_mask(kernelsize=10):
    def wrap(img):
        kernel = ee.Kernel.square(kernelsize, 'pixels')
        mask = img.mask()
        smooth = ee.Kernel.circle(radius=1)
        mask = mask.focal_max({'kernel': smooth, 'iterations': 2})
        mask = mask.focal_min({'kernel': kernel, 'iterations': 2})
        return img.updateMask(mask).copyProperties({'source': img})\
                  .set('system:time_start', img.get('system:time_start'))
    return wrap


import torch
from torch_geometric.loader import DataLoader

from loss import create_loss_fn
from model import FlexNet
from utils.data import flex_graph, \
    gen_rectangular_channel_matrix

def generate_data(n, k, batch_size):
    h_batch = gen_rectangular_channel_matrix(k, k, n, seed=11)
    datalist = flex_graph(h_batch)
    return DataLoader(datalist, batch_size=batch_size, shuffle=True)

def train(model, optimizer, loss_fn, dataset, k, path):
    model.train()
    for epoch in range(120):
        running_loss = 0.0
        for i, data in enumerate(dataset, 0):
            optimizer.zero_grad()
            out = model(data)
            loss = loss_fn(out, data.y)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            if i % 10 == 9:
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 10:.3f}')
                running_loss = 0.0
    torch.save(model.state_dict(), path)

def eval_model(path, k, noise_var):
    n = 10000
    model = FlexNet()
    model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
    model.eval()
    data_t = gen_rectangular_channel_matrix(k, k, n, seed=899)

    new_data = flex_graph(data_t)
    new_data = next(iter(DataLoader(new_data, batch_size=n, shuffle=False)))

    outs = model(new_data)
    p, t = outs
    t = torch.where(t >= 0.5, 1., 0.)
    p = torch.where(p >= 0.5, 1., 0.)
    rate = create_loss_fn(k, 1.)
    sum_rate = -rate((p, t), new_data.y).item()
    print('Test sum-rate:', sum_rate)
    return sum_rate

def create_and_train_model(n, k, batch_size, noise_var, path, lr=0.004, data=None, aggr='add'):
    dataset = data if (data is not None) else generate_data(n, k, batch_size)
    model = FlexNet(aggr)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = create_loss_fn(k, 1.)
    train(model, optimizer, loss_fn, dataset, k, path)

if __name__ == '__main__':
    path_ = 'flexible____.pth'
    # create_and_train_model(n=10000, k=20, batch_size=64, noise_var=1., path=path_)
    eval_model(path=path_, k=20, noise_var=1.)


import sys
from cloud.response import Response
from cloud.util import has_write_permission
from cloud.shortuuid import uuid
import base64

# Define the input output format of the function.
# This information is used when creating the *SDK*.
info = {
    'input_format': {
        'session_id': 'str',

        'parent_file_id': 'str?',
        'file_name': 'str?',
        'file_b64': 'str',

        'read_groups': 'list',
        'write_groups': 'list'
    },
    'output_format': {
        'success': 'bool',
        'message': 'str?',

        'file_id': 'str',
    }
}

def do(data, resource):
    body = {}
    params = data['params']
    user = data['user']

    user_id = user.get('id', None)

    parent_file_id = params.get('parent_file_id', None)
    file_name = params.get('file_name', uuid())
    file_b64 = params.get('file_b64')

    read_groups = params.get('read_groups', [])
    write_groups = params.get('write_groups', [])

    file_id = '{}'.format(uuid())
    parent_file_info = None

    file_size = sys.getsizeof(file_b64)

    if parent_file_id:
        parent_file_info = resource.db_get_item(parent_file_id)
        if parent_file_info:
            if has_write_permission(user, parent_file_info):
                file_name = parent_file_info.get('file_name', None)
                parent_file_info['next_file_id'] = file_id
                file_size += parent_file_info['file_size']
            else:
                body['success'] = False
                body['message'] = 'Permission denied'
                return Response(body)

    file_info = {
        'file_id': file_id,
        'parent_file_id': parent_file_id,
        'file_name': file_name,
        'owner': user_id,
        'file_size': file_size,
        'read_groups': read_groups,
        'write_groups': write_groups,
    }

    resource.db_put_item('files', file_info, file_id)
    if parent_file_id and parent_file_info:
        resource.db_update_item(parent_file_id, parent_file_info)

    file_b64 = file_b64.encode('utf-8')
    file_b64 = base64.b64decode(file_b64)
    resource.file_upload_bin(file_id, file_b64)

    body['success'] = True
    body['file_id'] = file_id
    return Response(body)


import re
import os
import io
import sys
import json
import http.client
import urllib
import urllib.request

PYPI_URL = "https://pypi.org/pypi/%s/json"
VERBOSE = False
COMPACT = False
INPLACE = False
YESMAN = False
SORTED = False

def err(*a):
    if VERBOSE:
        print(*a, file=sys.stderr)

class web:
    @staticmethod
    def get(url: str) -> bytes:
        err('url =', url)
        response: http.client.HTTPResponse = urllib.request.urlopen(url)
        body = response.read()
        assert response.status == 200, (response.status, url, body)
        return body

def get_name(line: str) -> str:
    m = re.match(r'^[\w\[\]-]*', line)
    assert m is not None
    return m.group()

def get_releases(name: str) -> dict:
    try:
        return json.loads(web.get(PYPI_URL % name).decode())['releases']
    except urllib.error.HTTPError as e:
        if e.code == 404:
            print(f'\u2715 No such package: {name!r}')
        else:
            print(f'\u2715 Could not retrieve info for package {name!r} (status code {e.code})')
        exit(1)
    except (KeyError, AssertionError) as e:
        print(f'\u2715 Could not retrieve info for package {name!r}')
        exit(1)

def get_latest(name: str) -> str:
    err('Fetching latest release number for', name)
    releases = get_releases(name)
    err('available releases:', list(releases.keys()))
    eligible_releases = [
        r for r in releases
        if re.match(r'^[0-9.]+$', r)
    ]
    latest_release = sorted(
        eligible_releases,
        key=lambda v: tuple(map(int, v.split('.')))
    )[-1]
    return latest_release

def round_up(filepath: str, file=sys.stdout, *, inplace: bool = False):

    with open(filepath) as f:
        lines = f.readlines()

    if inplace:
        file = open(filepath, mode='w')

    for line in map(str.strip, lines):
        if line.startswith('#') or not line:
            if not COMPACT:
                print(line, file=file)
            continue
        name = get_name(line)
        if line != name:
            print(line, file=file)
        else:
            vv = name + '==' + get_latest(name)
            print(vv, file=file)

    if inplace:
        print(f'\u2714 Wrote frozen pip requirements to {filepath}')

def yesno(question: str, default: str = 'y', skip: bool = False) -> bool:
    if skip is not None:
        return skip
    default = default.lower()
    ans = default.upper() + '/' + 'yn'.strip(default)
    return input(f'{question} ({ans})') in ['', 'y', 'Y']

def implicit():

    # detect file
    src = 'requirements.txt'
    if not os.path.exists(src):
        print('No pip requirements file found')
        exit(1)

    # write to requirements.txt if inplace
    dest = 'requirements.lock'
    if INPLACE:
        dest = src
    else:
        if os.path.exists(dest):
            YESMAN or print(f'\u2715 File {dest!r} already exists.', end=' ')
            cont = yesno('Replace ?', skip=YESMAN or None)
            # YESMAN and print()
            cont or exit(1)

    ss = io.StringIO()
    print('Fetching ...')
    round_up(src, ss)
    ss.seek(0)

    with open(dest, mode='w') as f:
        lines = ss.readlines()
        if SORTED:
            lines = sorted(lines)
        f.writelines(lines)

    print('\033[A', end='')
    print(f'\u2714 Wrote frozen pip requirements to {dest} ({len(lines)})')


import os
import sys
import gym
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.td3_actor import TD3Actor
from models.td3_critic import TD3Critic
from utils.tensors import get_tensor

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

class TD3Agent(object):
    def __init__(
            self,
            state_dim,
            action_dim,
            scale,
            model_path,
            actor_lr=0.0001,
            critic_lr=0.001,
            expl_noise=0.1,
            policy_noise=0.2,
            noise_clip=0.5,
            gamma=0.99,
            policy_freq=2,
            tau=0.005):
        self.scale = scale
        self.model_path = model_path

        # parameters
        self.expl_noise = expl_noise
        self.policy_noise = policy_noise
        self.noise_clip = noise_clip
        self.gamma = gamma
        self.policy_freq = policy_freq
        self.tau = tau

        self.actor = TD3Actor(state_dim, action_dim, scale=scale).to(device)
        self.actor_target = TD3Actor(state_dim, action_dim, scale=scale).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)

        self.critic1 = TD3Critic(state_dim, action_dim).to(device)
        self.critic2 = TD3Critic(state_dim, action_dim).to(device)
        self.critic1_target = TD3Critic(state_dim, action_dim).to(device)
        self.critic2_target = TD3Critic(state_dim, action_dim).to(device)

        self.critic1_optimizer = torch.optim.Adam(
                self.critic1.parameters(),
                lr=critic_lr, weight_decay=0.0001
        )
        self.critic2_optimizer = torch.optim.Adam(
                self.critic2.parameters(),
                lr=critic_lr, weight_decay=0.0001
        )

        #mean = torch.zeros(action_dim).to(device)
        #sigma = torch.ones(action_dim).to(device)
        #scale = self.expl_noise * self.actor.scale
        #self._exploration_noise = torch.distributions.Normal(mean, sigma)

        self._initialize_target_networks()

        self._initialized = False
        self.total_it = 0

    def _initialize_target_networks(self):
        self._update_target_network(self.critic1_target, self.critic1, 1.0)
        self._update_target_network(self.critic2_target, self.critic2, 1.0)
        self._update_target_network(self.actor_target, self.actor, 1.0)
        self._initialized = True

    def _update_target_network(self, target, origin, tau):
        for target_param, origin_param in zip(target.parameters(), origin.parameters()):
            target_param.data.copy_(tau * origin_param.data + (1.0 - tau) * target_param.data)

    def save(self, timestep=None):
        if not os.path.exists(os.path.dirname(self.model_path)):
            os.mkdir(os.path.dirname(self.model_path))

        if timestep:
            model_path = self.model_path + '_%i'%(timestep)
        else:
            model_path = self.model_path

        torch.save(self.actor.state_dict(), model_path+"_actor.h5")
        torch.save(self.actor_optimizer.state_dict(), model_path+"_actor_optimizer.h5")
        torch.save(self.critic1.state_dict(), model_path+"_critic1.h5")
        torch.save(self.critic2.state_dict(), model_path+"_critic2.h5")
        torch.save(self.critic1_optimizer.state_dict(), model_path+"_critic1_optimizer.h5")
        torch.save(self.critic2_optimizer.state_dict(), model_path+"_critic2_optimizer.h5")

    def load(self, timestep=-1):
        if timestep>0:
            model_path = self.model_path + '_%i'%(timestep)
        else:
            model_path = self.model_path

        self.actor.load_state_dict(torch.load(model_path+"_actor.h5"))
        self.actor_optimizer.load_state_dict(torch.load(model_path+"_actor_optimizer.h5"))
        self.critic1.load_state_dict(torch.load(model_path+"_critic1.h5"))
        self.critic2.load_state_dict(torch.load(model_path+"_critic2.h5"))
        self.critic1_optimizer.load_state_dict(torch.load(model_path+"_critic1_optimizer.h5"))
        self.critic2_optimizer.load_state_dict(torch.load(model_path+"_critic2_optimizer.h5"))

    def train(self, replay_buffer, iterations=1):
        if not self._initialized:
            self._initialize_target_networks()

        self.total_it += 1

        states, actions, n_states, rewards, not_done = replay_buffer.sample()

        with torch.no_grad():
            noise = (
                torch.randn_like(actions) * self.policy_noise
            ).clamp(-self.noise_clip, self.noise_clip)

            #n_actions = (
            #    self.actor_target(states) + noise
            #).clamp(-self.scale[0], self.scale[0])

            n_actions = self.actor_target(states) + noise
            n_actions = torch.min(n_actions,  self.actor.scale)
            n_actions = torch.max(n_actions, -self.actor.scale)

            target_Q1 = self.critic1_target(n_states, n_actions)
            target_Q2 = self.critic2_target(n_states, n_actions)
            target_Q = torch.min(target_Q1, target_Q2)
            target_Q = rewards + not_done * self.gamma * target_Q
            target_Q_detached = target_Q.detach()

        current_Q1 = self.critic1(states, actions)
        current_Q2 = self.critic2(states, actions)

        critic1_loss = F.mse_loss(current_Q1, target_Q_detached)
        critic2_loss = F.mse_loss(current_Q2, target_Q_detached)
        critic_loss = critic1_loss + critic2_loss

        self.critic1_optimizer.zero_grad()
        self.critic2_optimizer.zero_grad()
        critic_loss.backward()
        self.critic1_optimizer.step()
        self.critic2_optimizer.step()

        if self.total_it % self.policy_freq == 0:
            a = self.actor(states)
            Q1 = self.critic1(states, a)
            actor_loss = -Q1.mean()

            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            self._update_target_network(self.critic1_target, self.critic1, self.tau)
            self._update_target_network(self.critic2_target, self.critic2, self.tau)
            self._update_target_network(self.actor_target, self.actor, self.tau)

            return {'actor_loss': actor_loss, 'critic_loss': critic_loss}

        return {'critic_loss': critic_loss}

    def policy(self, state, to_numpy=True):
        state = get_tensor(state)
        action = self.actor(state)

        if to_numpy:
            return action.cpu().data.numpy().squeeze()

        return action.squeeze()

    def policy_with_noise(self, state, to_numpy=True):
        state = get_tensor(state)
        action = self.actor(state)

        action = action + self._sample_exploration_noise(action) # actions.shape
        action = torch.min(action,  self.actor.scale)
        action = torch.max(action, -self.actor.scale)

        if to_numpy:
            return action.cpu().data.numpy().squeeze()

        return action.squeeze()

    def _sample_exploration_noise(self, actions): # shape
        scale = self.expl_noise * self.actor.scale
        mean = torch.zeros(actions.size()).to(device)
        var = torch.ones(actions.size()).to(device)
        return torch.normal(mean, scale*var)

        # return self._exploration_noise.sample(sample_shape=shape)

    def evaluate_policy(self, env, eval_episodes=10, render=False, save_video=False, sleep=-1):
        if save_video:
            from OpenGL import GL
            env = gym.wrappers.Monitor(env, directory='video',
                                    write_upon_reset=True, force=True, resume=True, mode='evaluation')
            render = False

        rewards = []
        for e in range(eval_episodes):
            s = env.reset()
            done = False
            steps = 1
            reward_episode_sum = 0

            while not done:
                if render:
                    env.render()
                if sleep>0:
                    time.sleep(sleep)
                a = self.policy(s)
                n_s, r, done, _ = env.step(a)

                s = n_s
                reward_episode_sum += r
                steps += 1
            else:
                print("Rewards in Episode %i: %.2f"%(e, reward_episode_sum/steps))
                rewards.append(reward_episode_sum)

        return np.array(rewards)


import os
import re
import time

from dotenv import main

from selenium.webdriver.common.by import By

from utils.directory import create_directories
from utils.files import write, download, open_list
from utils.helper import update_link_files, get_links
from utils.translate import page

main.load_dotenv()

website_link = os.getenv('WEBSITE_LINK')
parent_folder = os.getenv('PARENT_FOLDER')

def html(driver, files_only=False, scripts_only=False):
    print('START:: HTML')

    driver.get(website_link)
    assert "Class Central" in driver.title

    file_link_details = []

    include_html = (not files_only and not scripts_only)

    if files_only or include_html:
        file_link_details.append({
            'path': 'files/scraper_files.json',
            'in_filter': ['.png', '.svg', '.webmanifest', '.woff2'],
            'not_in_filter': ['catalog-iframe']
        })

    if scripts_only or include_html:
        file_link_details.append({
            'path': 'files/scraper_js.json',
            'in_filter': ['.js'],
            'not_in_filter': ['catalog-iframe']
        })

    # update link files to be use for script and files downloading and scraping
    if len(file_link_details) != 0:
        index_links = get_links(driver)
        update_link_files(index_links, file_link_details)

    if include_html:
        html = driver.page_source
        index_file = parent_folder + '/index.html'

        language = os.getenv('TRANSLATE_LANG')

        # translate page to hindi
        translated_html = page(html, language, convert_links=True)

        # create index file
        write(index_file, 'w', '<!DOCTYPE html>')
        write(index_file, 'a', translated_html, 'utf-8')

    # scrape links
    links = driver.find_elements(By.TAG_NAME, 'a')
    link_details_list = []

    retry = 0
    trigger = True

    while trigger:
        try:
            for link in links:
                href = link.get_attribute('href')
                path = href.replace(website_link, '')

                is_exist = len([x for x in link_details_list if x['path'] == path]) != 0

                if website_link in href and path and not is_exist:
                    link_details_list.append({
                        'href': href,
                        'path': path
                    })

            link_details = []

            if len(file_link_details) != 0 and not include_html:
                link_details = link_details_list
            else:
                link_details = [x for x in link_details_list if not os.path.exists(parent_folder + '/' + x['path'])]

            counter = 0

            for detail in link_details:
                href = detail['href']
                path = detail['path']

                counter += 1;
                print('SCRAPING::', href, path, counter, len(link_details))

                driver.get(href)
                time.sleep(3)

                # update link files to be use for script and files downloading and scraping
                if len(file_link_details) != 0:
                    inner_links = get_links(driver)
                    update_link_files(inner_links, file_link_details)

                if include_html:
                    relative_path = parent_folder + '/' + path;
                    inner_index_file = relative_path + '/' + 'index.html'
                    inner_html = driver.page_source

                    # translate page to hindi
                    translated_html = page(inner_html, language, convert_links=True)

                    create_directories(relative_path)

                    # create index file
                    write(inner_index_file, 'w', '<!DOCTYPE html>')
                    write(inner_index_file, 'a', translated_html, "utf-8" )

            trigger = False
        except Exception as error:
            retry += 1
            trigger = retry <= 3

            print('RETRY::', retry)

            if not trigger:
                print('RETRY::ERROR', error)

    print('FINISHED:: HTML')

def files():
    print('START:: FILES')

    files_link = open_list('files/scraper_files.json')

    files_counter = 0

    for link in files_link:
        folders = link.replace(website_link, '').split('/')
        folders.pop()
        folders.insert(0, parent_folder)

        files_counter += 1

        create_directories('/'.join(folders))

        print('DOWNLOAD::', files_counter, len(files_link))
        response = download(link)

        write(parent_folder + '/' + link.replace(website_link, ''), 'wb', response.content)

    print('FINISHED:: FILES')

def scripts(driver):
    print('START:: SCRIPTS')

    script_links = open_list('files/scraper_js.json')

    script_counter = 0

    for script in script_links:
        href = script

        folders = script.replace(website_link, '').split('/')
        folders.pop()
        folders.insert(0, parent_folder)

        create_directories('/'.join(folders))

        script_counter += 1
        print('SCRAPING::', href, script_counter, len(script_links))

        driver.get(href)
        time.sleep(3)

        script_element = driver.find_element(By.TAG_NAME, 'pre')
        jsscript = script_element.get_attribute('innerText')

        write(parent_folder + '/' + script.split('?')[0].replace(website_link, ''), 'w', jsscript, 'utf-8')

    print('FINISHED:: SCRIPTS')

def link_rel(driver):
    print('START:: LINK_REL')

    driver.get(website_link)
    assert "Class Central" in driver.title

    rel_list = ['manifest', 'apple-touch-icon', 'icon', 'mask-icon', 'search']

    for rel in rel_list:
        script_element = driver.find_element(By.CSS_SELECTOR, 'link[rel="'+ rel +'"]')

        if script_element:
            href = script_element.get_attribute('href')

            print('DOWNLOADING::', rel, href)

            response = download(href)

            write(parent_folder + '/' + href.replace(website_link, ''), 'wb', response.content)

    print('FINISHED:: LINK_REL')


import importlib.util
import json
import logging
import os
import sys
from pathlib import Path

from starlette.applications import Starlette
from starlette.responses import JSONResponse, PlainTextResponse
from starlette.routing import Route

import traceback
import typer
import uvicorn

from .serializer_deserializer import decode

logging.basicConfig(format="%(levelname)s:     %(message)s", datefmt="", level=logging.INFO)
logger = logging.getLogger("uvicorn.access")

async def some_startup_task():
    global inference_handler

    _handler_path = Path(os.environ.get("HANDLER_PATH", "handler.py"))
    module_name = f"{_handler_path.stem}.EndpointHandler"

    spec = importlib.util.spec_from_file_location(module_name, _handler_path)
    # add the whole directory to path for submodlues
    sys.path.insert(0, str(_handler_path.parents[0]))
    # import custom handler
    handler = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = handler
    spec.loader.exec_module(handler)
    # init custom handler with model_dir
    inference_handler = handler.EndpointHandler(str(_handler_path.parents[0]))
    # logger.info(f"Loaded EndpointHandler from {str(_handler_path)} file")

async def predict(request):
    try:
        content_type = request.headers.get("content-Type", None)
        deserialized_body = decode(await request.body(), content_type)
        # checks if input schema is correct
        if "inputs" not in deserialized_body:
            raise ValueError(f"Body needs to provide a inputs key, recieved: {json.dumps(deserialized_body)}")

        # run handler call
        pred = inference_handler(deserialized_body)
        return JSONResponse(pred)
    except Exception as e:
        print(traceback.format_exc())
        return JSONResponse({"error": str(e)}, status_code=400)

async def health(request):
    return PlainTextResponse("Ok")

app = app = Starlette(
    debug=True,
    routes=[
        Route("/", health, methods=["GET"]),
        Route("/health", health, methods=["GET"]),
        Route("/", predict, methods=["POST"]),
        Route("/predict", predict, methods=["POST"]),
    ],
    on_startup=[some_startup_task],
)

def emulate(handler: str = "handler.py", port: int = 5000):
    os.environ["HANDLER_PATH"] = handler
    logger.info(f"Loading handler from {handler} file")
    uvicorn.run(app, port=port, host="0.0.0.0", log_level="info")

if __name__ == "__main__":
    typer.run(emulate)


import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms

import os
import numpy as np
from heapq import nlargest

from utils import PCamDataset

SAVE_MODEL_PATH = "/save/train_resnet_sequential_unfreeze_best20200604T175157"

ROOT_DIR = os.getcwd().replace("\\", "/")

# DATASET
BATCH_SIZE = 8 #  To train on small dataset
EVALUATE_EVERY = 4 #  To train on small dataset
SIZE_TRAIN_DATASET, SIZE_VAL_DATASET = 16, 128 #  To train on small dataset
#BATCH_SIZE = 256
# EVALUATE_EVERY = 50000
#SIZE_TRAIN_DATASET, SIZE_VAL_DATASET = None, None

# DEVICE SETUP
USE_GPU = True

if USE_GPU and torch.cuda.is_available():
    device = torch.device('cuda')
    dtype = "torch.cuda.FloatTensor"
else:
    device = torch.device('cpu')
    dtype = "torch.FloatTensor"

print('using device:', device)

def compute_saliency_maps(X, y, model):
    """
    Compute a class saliency map using the model for images X and labels y.

    Input:
    - X: Input images; Tensor of shape (N, 3, H, W)
    - y: Labels for X; LongTensor of shape (N,)
    - model: A pretrained CNN that will be used to compute the saliency map.

    Returns:
    - saliency: A Tensor of shape (N, H, W) giving the saliency maps for the input
    images.
    """
    # Make sure the model is in "test" mode
    model.eval()

    # Make input tensor require gradient
    X.requires_grad_()

    # Compute saliency
    N, _, H, W = X.shape
    scores = model.forward(X)
    loss = scores.sum(axis=0)
    loss.backward()
    X_grad = X.grad
    saliency, _ = X_grad.abs().max(axis=1)

    return saliency

def main():
    data_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    data_detransform = transforms.Compose([
        transforms.Normalize(mean=[0, 0, 0], std=[1/0.229, 1/0.224, 1/0.225]),
        transforms.Normalize(mean=[-0.485, -0.456, -0.406], std=[1, 1, 1]),
    ])

    train_dataset = PCamDataset(csv_file="train_labels.csv",
                                root_dir=ROOT_DIR+"/data/",
                                transform=data_transform)
    val_dataset = PCamDataset(csv_file="dev_labels.csv",
                              root_dir=ROOT_DIR+"/data/",
                              transform=data_transform)

    if SIZE_TRAIN_DATASET is not None:
        train_dataset = torch.utils.data.Subset(dataset=train_dataset,
                        indices=np.random.choice(
                            a=np.arange(len(train_dataset)),
                            size=SIZE_TRAIN_DATASET,
                            replace=False))
    if SIZE_VAL_DATASET is not None:
        val_dataset = torch.utils.data.Subset(dataset=val_dataset,
                        indices=np.random.choice(
                            a=np.arange(len(val_dataset)),
                            size=SIZE_VAL_DATASET,
                            replace=False))

    loader_train = DataLoader(train_dataset,
                          batch_size=BATCH_SIZE,
                          shuffle=True)

    loader_val = DataLoader(val_dataset,
                        batch_size=BATCH_SIZE,
                        shuffle=True)

    model = torch.load(ROOT_DIR + SAVE_MODEL_PATH.replace('.', ''), map_location=torch.device('cpu'))
    model.eval()

    with torch.no_grad():

        y_prob = torch.zeros(val_dataset.__len__())
        x_val = torch.zeros((val_dataset.__len__(), 3, 96, 96))
        y_val = torch.zeros(val_dataset.__len__())
        for t, (x, y) in enumerate(loader_val):
            x = x.to(device=device)  # move to device, e.g. GPU
            y = y.to(device=device)
            scores = model(x)
            y_prob[t*BATCH_SIZE: min((t+1)*BATCH_SIZE, val_dataset.__len__())] = torch.sigmoid(scores).flatten()
            y_val[t*BATCH_SIZE: min((t+1)*BATCH_SIZE, val_dataset.__len__())] = y.flatten()
            x_val[t*BATCH_SIZE: min((t+1)*BATCH_SIZE, val_dataset.__len__()), :, :, :] = x

    # Saliency maps

    # True Positive
    mask = y_val
    n = 3
    tuples = nlargest(n, enumerate(y_prob*mask), key=lambda x: x[1])

    X = torch.stack([x_val[idx] for idx, prob in tuples])
    y = torch.stack([y_val[idx] for idx, prob in tuples])
    saliency = compute_saliency_maps(X, y, model)

    for i in range(n):
        saliency_im = transforms.ToPILImage(mode=None)(saliency[i].reshape(1, 96, 96))
        input_im = transforms.ToPILImage(mode=None)(data_detransform(x[i]))
        input_im.save("./save_image/TP_"+str(i)+"_original_image.png")
        saliency_im.save("./save_image/TP_"+str(i)+"_saliency_map.png")

    # False Negative
    mask = y_val
    n = 3
    tuples = nlargest(n, enumerate((1-y_prob)*mask), key=lambda x: x[1])

    X = torch.stack([x_val[idx] for idx, prob in tuples])
    y = torch.stack([y_val[idx] for idx, prob in tuples])
    saliency = compute_saliency_maps(X, y, model)

    for i in range(n):
        saliency_im = transforms.ToPILImage(mode=None)(saliency[i].reshape(1, 96, 96))
        input_im = transforms.ToPILImage(mode=None)(data_detransform(x[i]))
        input_im.save("./save_image/FN_"+str(i)+"_original_image.png")
        saliency_im.save("./save_image/FN_"+str(i)+"_saliency_map.png")

if __name__ == "__main__":
    pass
    main()


#%%
def charToPriority(char):
    if 'a' <= char <= 'z':
        return ord(char) - ord('a') + 1
    elif 'A' <= char <= 'Z':
        return ord(char) - ord('A') + 27
    else:
        return 0
#%%
with open("input", "r") as f:
    rows = [x.strip() for x in f.readlines()]

#%%
total = 0
for r in rows:
    total += charToPriority((set(r[:len(r)//2]) & set(r[len(r)//2:])).pop())
print(total)
# %%
from itertools import zip_longest
from functools import reduce
#%%
total = 0
N = 3
for rs in zip_longest(*( [iter(rows)]*N)):
    intersection = reduce(lambda x,y: set(x) & set(y), rs)
    total += charToPriority(intersection.pop())
print(total)
# %%


import os
import sys
import time
import logging
import argparse
import itertools
import matplotlib; matplotlib.use('Agg')
import multiprocessing
import pandas as pd
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

from operator import add
from functools import partial
from collections import Counter
from matplotlib import font_manager
from GetConfig import getConfig

config = getConfig()

'''
Description:
This module is used for network analysis.
Network analysis is an efficient method to show relationship of haplotypes.
Y haplotype data is required for analysis.
It is recommended to use Y-STR haplotype under a specific haplogroup.
Network analysis can show the evolutionary relationship between haplotypes.
This command can output network figure, and fdi file.
'''

def net_parser():

    parser = argparse.ArgumentParser('net', description='(c) Y-LineageTracker: Network analysis')
    # function used for network analysis
    parser.add_argument('net',
                        help='Perform network analysis from Y haplotype data.')
    # required, input file, alignment sequence or matrix
    input = parser.add_mutually_exclusive_group(required=True)
    input.add_argument('--seq',
                        type=str,
                        action='store',
                        help='seq: Haplotype data in sequence alignment format')
    input.add_argument('--matrix',
                        type=str,
                        action='store',
                        help='martix: Y-STR haplotype data in matrix format')
    # optional, format of sequence alignment
    parser.add_argument('--seq-format',
                        required=False,
                        type=str,
                        dest='format',
                        action='store',
                        choices=['fasta', 'phylip', 'nexus', 'meg', 'vcf'],
                        help='--seq-format: The format of sequence file. This option is only required for the sequence alignment file. ')
    # optional, population file
    parser.add_argument('-p', '--population',
                        required=False,
                        type=str,
                        action='store',
                        help='population: A file containing sample ID and population information of each individual.')
    # optional, type of tree
    parser.add_argument('--tree-type',
                        required=False,
                        type=str,
                        dest='treetype',
                        action='store',
                        default='mjn',
                        choices=['mjn', 'msn'],
                        help='tree-type: The algorithm used to construct a network tree.')
    # optional, how to code missing site
    parser.add_argument('--gap',
                        required=False,
                        type=str,
                        action='store',
                        default='allele',
                        choices=['allele', 'missing'],
                        help='gap: The gap coding method. You can choose to code gaps as a new type of allele or missing data')
    # optional, the search for new median nodes
    parser.add_argument('--search',
                        required=False,
                        type=int,
                        action='store',
                        default=0,
                        help='search: How the search for new median nodes is performed: the larger this parameter, the wider the search. ')
    # optional, output fdi file
    parser.add_argument('--fdi',
                        required=False,
                        action='store_true',
                        help='fdi: Output fdi file that can be used as the input file in Network software, to show network plot or estimate TMRCA directly.')
    # optional, filter sites by misisng rate value
    parser.add_argument('--filter',
                        required=False,
                        type=float,
                        action='store',
                        default=0.7,
                        help='filter: The cutoff value of the site to be filtered, sites with a missing rate value greater than this value will be filtered')
    # optional, how to show alternative link in figure
    parser.add_argument('--alt',
                        required=False,
                        type=str,
                        action='store',
                        default='hide',
                        choices=['hide', 'show', 'dashed'],
                        help='alt: Whether to show alternative links between nodes in the network figure.')
    # optional, size of median nodes
    parser.add_argument('--mediansize',
                        required=False,
                        type=float,
                        action='store',
                        default=0.0,
                        help='mediansize: The median node size in the median-joining network compared to base size in the network figure. The default is 0.')
    # optional, keep all nodes in same size
    parser.add_argument('--euqal-size',
                        required=False,
                        dest='euqalsize',
                        action='store_true',
                        help='euqal-size: Keep all nodes the same size in the network figure.')
    # optional, the prefix of output
    parser.add_argument('-o', '--output',
                        required=False,
                        type=str,
                        action='store',
                        help='output: The prefix of output files. ')

    args = parser.parse_args()

    return args

# check missing rate is filter argument is used
def check_misisng_rate(missing_rate):

    if missing_rate < 0 or missing_rate> 1:
        print('[Y-LineageTracker] [Error] Cutoff of missing rate should be a float between 0 and 1')
        sys.exit()

# print program information and write to log file
def set_log(args_log, log_file):

    logger = logging.getLogger()
    logger.setLevel(level=logging.INFO)

    handler = logging.FileHandler(log_file, mode='w')
    handler.setLevel(logging.INFO)
    formatter = logging.Formatter('[%(asctime)s] - [%(levelname)s]: %(message)s')
    handler.setFormatter(formatter)

    console = logging.StreamHandler()
    console.setLevel(logging.INFO)

    logger.addHandler(handler)
    logger.addHandler(console)

    log_info = ['[Y-LineageTracker] [Network]',
                '[Y-LineageTracker] Run Date: ' + time.asctime(time.localtime(time.time())),
                '[Y-LineageTracker] Haplotype File: %s' % args_log.seq,
                '[Y-LineageTracker] Haplotype File Format: %s' % args_log.format,
                '[Y-LineageTracker] Gap euqal to %s' % args_log.gap,
                '[Y-LineageTracker] Searching parameter: %d' % args_log.search,
                '[Y-LineageTracker] Missing rate of allele to be filtered: %s' % args_log.filter,
                '[Y-LineageTracker] Alternative link type : %s' % args_log.alt,
                '[Y-LineageTracker] Median node size: %s' % args_log.mediansize]

    if args_log.euqalsize:
        log_info.append('[Y-LineageTracker] Keep all node the same size')
    else:
        log_info.append('[Y-LineageTracker] Plot node according to sample size')

    print('\n')
    for i in log_info:
        logger.info(i)

class HaploNetwork(object):
    '''
    This class is used for network analysis.
    Main functions of this class include:
    1. construct median-joining network from haplotypes
    2. plot network by networkx package
    3. write network to fdi file
    4. output network info
    '''

    def __init__(self, gap, path):

        self.logger = logging.getLogger()
        self.skip_letters = config.get('Statistics', 'SkipLetters').split(',')
        self.gap = gap
        self.path = path
        self.net_info = open(self.path+'.info.net', 'w')

    # convert STR numbers to 0-1 code haplotype
    def _convert_allel_to_binary(self, haplotype_data, filter):

        # create a empty dataframe to save haplotypes
        sequence_df = pd.DataFrame(index=haplotype_data.index,
                                   columns=['Haplotype'])
        sequence_df['Haplotype'] = ''

        STR_code_pos = {}
        start_pos = 0
        hap_length = 0
        allel_num = haplotype_data.index.size
        for STR in haplotype_data.columns.tolist():
            str_allels = haplotype_data[STR].astype(str).tolist()
            # filter STR according to filter parameter
            if sum([str_allels.count(i) for i in self.skip_letters])/allel_num >= filter:
                continue
            if len(set(set(str_allels)&set(self.skip_letters))) > 1:
                str_allels = ['.' if i in self.skip_letters else i for i in str_allels]
            repeats_type = sorted(list(set(str_allels)))
            # convert STR repeat numbers to 0-1 code
            repeats_type_num = len(repeats_type)
            max_length = len('{0:b}'.format(repeats_type_num-1))
            binary_set = [] # creat a empty list to save 0-1 code
            for i in range(repeats_type_num):
                binary_num = '{0:b}'.format(i)
                binary_num_length = len(binary_num)
                if binary_num_length < max_length:
                    binary_num = (max_length-binary_num_length)*'0'+binary_num
                binary_set.append(binary_num)
            if self.gap == 'missing':
                replaced_allels = ['-'*max_length if i == '.' else binary_set[repeats_type.index(i)] for i in str_allels]
            else:
                replaced_allels = [binary_set[repeats_type.index(i)] for i in str_allels]
            sequence_df['Haplotype'] = list(map(add, sequence_df['Haplotype'].tolist(), replaced_allels))
            for pos in range(start_pos, start_pos+max_length):
                STR_code_pos[pos] = STR
            start_pos += max_length
            hap_length += 1

        self.net_info.write('#HapLength:%d\n\n' % hap_length)

        return sequence_df, STR_code_pos

    # calculate frequency of all kinds of haplotypes and the proportion in each population
    def _calculate_distance_matrix(self, sequence_df):

        # get haplotypes and name for them
        haplotypes = sorted(list(set(sequence_df['Haplotype'])))
        sp_sequence_df = pd.DataFrame(index=range(len(haplotypes)), columns=['Haplotype'])
        sp_sequence_df['Haplotype'] = haplotypes

        # calculate distance between two haplotypes
        distance_matrix = pd.DataFrame(index=range(len(haplotypes)), columns=range(len(haplotypes)))
        hap_num = 0
        for hap1_name, hap1 in enumerate(haplotypes):
            hap_num += 1
            for hap2_name, hap2 in enumerate(haplotypes[hap_num:]):
                if self.gap == 'missing':
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2) if l1 not in self.skip_letters and l2 not in self.skip_letters)
                    if distance == 0:
                        distance = 1
                else:
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2))
                hap2_name = hap2_name + hap_num
                distance_matrix.at[hap2_name, hap1_name] = distance

        self.net_info.write('#HapSamples:\n')

        sample_sp_code = {}
        for i in haplotypes:
            idx = sequence_df['Haplotype'].tolist().index(i)
            sample = sequence_df.index[idx]
            sample_sp_code[sample] = i
            all_samples = ','.join(sequence_df[sequence_df['Haplotype']==i].index.tolist())
            self.net_info.write('%d\t%s\n' %(idx, all_samples))

        self.net_info.write('\n')

        return distance_matrix, sp_sequence_df, sample_sp_code

    # get msn tree
    def _calculate_minimum_spanding_tree(self, n, distance_matrix, return_pruned=True):

        if n < 2:
            self.logger.error('[Y-LineageTracker] Program stopped since there are less than 2 types of haplotypes')
            sys.exit()

        # create a empty edge matrix to save hapotype and their distance
        edge_num = n-1
        edge_matrix = np.empty(shape=(edge_num, 3))
        edge_matrix[:] = np.nan

        # get the distance from small to large between of two haplotypes
        forest = np.linspace(0, n-1, n)

        add_num = 0
        while add_num < edge_num:
            ci, ri = np.unravel_index(np.nanargmin(distance_matrix.T.values), distance_matrix.T.shape)
            row_record = forest[ri]
            col_record = forest[ci]
            if row_record != col_record:
                edge_matrix[add_num] = [ci, ri, distance_matrix.at[ri, ci]] # col hap, row hap, step
                forest[forest==col_record] = row_record
                add_num += 1
            distance_matrix.at[ri, ci] = np.nan

        if return_pruned:
            return edge_matrix, distance_matrix
        else:
            return edge_matrix

    # get median haplotypes
    @staticmethod
    def calculate_median_vector(triplet, sp_sequence_df, max_n_cost, epsilon, n, gap, new_link_data, shared_median_haps, initial_lambda, n_cost, lock):

        median_haps = pd.DataFrame(columns=['Haplotype'])
        temp_link_data = np.empty(shape=(0, 3))
        median_num = n

        for item in triplet:
            hap_in_order = item
            # check whether two haplotypes can be linked
            condition1 = any(hap_in_order[0]==np.array(new_link_data)[:, 0]) and any(hap_in_order[1]==np.array(new_link_data)[:, 1])
            condition2 = any(hap_in_order[0]==np.array(new_link_data)[:, 0]) and any(hap_in_order[2]==np.array(new_link_data)[:, 1])
            condition3 = any(hap_in_order[1]==np.array(new_link_data)[:, 0]) and any(hap_in_order[2]==np.array(new_link_data)[:, 1])
            if (condition1+condition2+condition3) < 2:
                continue
            # get alleles that are variable
            sub_seq = sp_sequence_df.loc[list(hap_in_order)]
            # get the median haplotype
            sub_seq_length = len(sub_seq.at[hap_in_order[0], 'Haplotype'])
            median_seq = ''.join([Counter(sub_seq['Haplotype'].map(lambda x: x[i])).most_common(1)[0][0] for i in range(sub_seq_length)])
            # check whether median sequence exist
            if median_seq in sp_sequence_df['Haplotype'].tolist():
                continue
            if not median_haps.empty:
                if median_seq in shared_median_haps:
                    continue
            # calculate the distance between median sequence and haplotypes
            if gap == 'missing':
                median_distance = [sum (l1 != l2 for l1, l2 in zip(i, median_seq) if l1 not in self.skip_letters and l2 not in self.skip_letters) for i in sub_seq['Haplotype']]
                if 0 in median_distance:
                    median_distance = [1 if dist==0 else dist for dist in median_distance]
            else:
                median_distance = [sum (l1 != l2 for l1, l2 in zip(i, median_seq)) for i in sub_seq['Haplotype']]
            connection_cost = sum(median_distance)
            # update n cost
            lock.acquire()
            n_cost.value += 1
            if n_cost.value > max_n_cost:
                print('[Y-LineageTracker] Program stopped since there are too many connnections')
                sys.exit()
            if connection_cost < initial_lambda.value:
                # update initial lambda
                initial_lambda.value = connection_cost
            if connection_cost <= initial_lambda.value + epsilon:
                hap_to_add1 = hap_in_order
                hap_to_add2 = [median_num]*3
                step_to_add = median_distance
                link_data_to_add = np.array([hap_to_add1, hap_to_add2, step_to_add]).T
                temp_link_data = np.concatenate([temp_link_data, link_data_to_add])
                median_haps.loc[median_num] = median_seq
                median_num += 1
                if median_seq not in shared_median_haps:
                    shared_median_haps.append(median_seq)
            lock.release()

        return median_haps, temp_link_data

    # construct network from distance matrix
    def _calculate_network(self, distance_matrix, sp_sequence_df, tree_type, epsilon):

        # initialization
        n = sp_sequence_df.index.size
        haplotype_length = len(sp_sequence_df.at[0, 'Haplotype'])
        median_num = n
        MST, pruned_distance_matrix = self._calculate_minimum_spanding_tree(n ,distance_matrix)

        if n < 3 or tree_type == 'msn':
            self.logger.warning('There are less than 3 types of haplotypes')
            alt_data = np.empty(shape=(0, 3))
            return MST, alt_data

        MST_step = MST[:, 2]
        lower_limit = min(MST_step)
        upper_limit = max(MST_step) + epsilon
        new_link_data = []
        for i in pruned_distance_matrix.columns:
            for j in pruned_distance_matrix.index:
                step = pruned_distance_matrix.at[j, i]
                if (not np.isnan(step)) and step > lower_limit and step <= upper_limit:
                    new_link_data.append([i, j, step])
        if new_link_data:
            new_link_data = np.concatenate([MST, np.array(new_link_data)])
        else:
            new_link_data = MST

        start_n_cost = 0
        start_initial_lambda = float('inf')
        tripleted = []

        while True:

            # get triplet for calculation
            triplet_new = list(itertools.combinations(range(n), 3))
            triplet = [i for i in triplet_new if i not in tripleted]
            tripleted.extend(triplet)

            process_num = 10
            pool = multiprocessing.Pool(processes=process_num)

            # global variables to be shared
            manager = multiprocessing.Manager()
            n_cost = manager.Value('i', start_n_cost)
            initial_lambda  = manager.Value('f', start_initial_lambda)
            shared_new_link_data = manager.list(new_link_data)
            shared_median_haps = manager.list([])
            LOCK = manager.Lock()

            # split triplet to sections
            triplet_interval = np.linspace(0, len(triplet), process_num, dtype=int)
            #triplet_sections = [triplet[triplet_interval[n]: triplet_interval[n+1]] if n == 0 else triplet[(triplet_interval[n]+1): triplet_interval[n+1]] for n in range(process_num-1)]
            triplet_sections = [triplet[triplet_interval[n]: triplet_interval[n+1]] for n in range(process_num-1)]

            # multiprocess triplet data
            pt = partial(self.calculate_median_vector, # function
                         sp_sequence_df=sp_sequence_df, # variable for read
                         max_n_cost=40000, # variable for read
                         epsilon=epsilon, # variable for read
                         n=n, # variable for read
                         gap=self.gap, # variable for read
                         new_link_data=new_link_data, # variable for writing
                         shared_median_haps=shared_median_haps, # variable for writing
                         initial_lambda=initial_lambda, # variable for writing
                         n_cost=n_cost, # variable for writing
                         lock=LOCK) # process lock
            median_haps_list, temp_link_data_list = zip(*pool.map(pt, triplet_sections))

            pool.close()
            pool.join()

            # intergrate data
            median_haps = pd.concat(median_haps_list)
            temp_link_data = np.concatenate(temp_link_data_list)
            if median_haps.empty:
                self.logger.info('[Y-LineageTracker] Get all median haplotypes')
                break
            else:
                # code for median vectors
                start_num = sp_sequence_df.index.size
                end_num = start_num + median_haps.index.size
                median_haps.index = range(start_num, end_num)
                temp_link_data[:, 2] = [val for val in range(start_num, end_num) for _ in (0, 1, 2)]
                # add to haplotype information
                sp_sequence_df = pd.concat([sp_sequence_df, median_haps])
                new_link_data = np.concatenate([new_link_data, temp_link_data])
                # update for iteration
                start_n_cost = n_cost.value
                start_initial_lambda = initial_lambda.value

        sp_sequence_df = sp_sequence_df.drop_duplicates().reset_index()

        self.logger.info('[Y-LineageTracker] Calculating distance matrix...')
        # calculate and add median haplotype to distance matrix
        hap_num = 0
        for hap1_name, hap1 in zip(sp_sequence_df.index, sp_sequence_df['Haplotype']):
            hap_num += 1
            for hap2_name, hap2 in zip(sp_sequence_df.index[hap_num:], sp_sequence_df['Haplotype'][hap_num:]):
                if self.gap == 'missing':
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2) if l1 not in self.skip_letters and l2 not in self.skip_letters)
                    if distance == 0:
                        distance = 1
                else:
                    distance = sum(l1 != l2 for l1, l2 in zip(hap1, hap2))
                distance_matrix.at[hap2_name, hap1_name] = distance

        N = sp_sequence_df.index.size
        dnet = distance_matrix.copy().applymap(lambda x: 0 if not np.isnan(x) else x)
        forest = np.linspace(0, N-1, N)
        network_data = np.empty(shape=(0, 3))
        alt_data = np.empty(shape=(0, 3))
        steps = sorted(list(set(filter(lambda x: x==x, distance_matrix.to_numpy().reshape(-1)))))

        edge_num = 0
        add_new_data = False

        self.logger.info('[Y-LineageTracker] Constructing relationship between nodes')
        while len(set(forest)) > 1:

            step_num = steps[edge_num]
            step_df = distance_matrix[distance_matrix==step_num]
            step_count = step_df.notnull().sum().sum()
            for i in range(step_count):
                ci, ri = np.unravel_index(np.nanargmin(step_df.T.values), step_df.T.shape)
                step_df.at[ri, ci] = np.nan
                row_record = forest[ri]
                col_record = forest[ci]
                if row_record != col_record:

                    for col in np.where(forest==col_record)[0]:
                        if ci > col:
                            num_col = dnet.at[ci, col]
                        elif ci < col:
                            num_col = dnet.at[col, ci]
                        else:
                            num_col = 0
                        for row in np.where(forest==row_record)[0]:
                            if ri > row:
                                num_row = dnet.at[ri, row]
                            elif ri < row:
                                num_row = dnet.at[row, ri]
                            else:
                                num_row = 0
                            dnet_num = num_col + num_row + step_num
                            dnet.at[row, col] = dnet_num
                    forest[forest==col_record] = row_record
                    add_new_data = True
                else:
                    if step_num < dnet.at[ri, ci]:
                        add_new_data = True
                        # update dnet
                        cols = np.append(network_data[network_data[:, 1]==ci, 0], network_data[network_data[:, 0]==ci, 1])
                        rows = np.append(network_data[network_data[:, 1]==ri, 0], network_data[network_data[:, 0]==ri, 1])
                        for col in np.append(ci, cols).astype(int):
                            if ci > col:
                                num_col = dnet.at[ci, col]
                            elif ci < col:
                                num_col = dnet.at[col, ci]
                            else:
                                num_col = 0
                            for row in np.append(ri, rows).astype(int):
                                if col == row:
                                    continue
                                else:
                                    if ri > row:
                                        num_row = dnet.at[ri, row]
                                    elif ri < row:
                                        num_row = dnet.at[row, ri]
                                    else:
                                        num_row = 0
                                    dnet_num = num_col + num_row + step_num
                                    if row > col:
                                        if dnet.at[row, col] > dnet_num:
                                            dnet.at[row, col] = dnet_num
                                    elif row < col:
                                        if dnet.at[col, row] > dnet_num:
                                            dnet.at[col, row] = dnet_num
                if add_new_data:
                    network_to_add = np.array([[ci, ri, step_num]])
                    network_data = np.concatenate([network_data, network_to_add])
                    add_new_data = False
            edge_num += 1
            if edge_num > 5000:
                self.logger.error('Program stopped since there are too many connnections')
                sys.exit()

        network_data = network_data.astype(int)
        net_MST = self._calculate_minimum_spanding_tree(N ,distance_matrix, False)
        net_MST = net_MST.astype(int)

        if len(network_data) > len(net_MST):
            alt_data = np.array([i for i in network_data if i[:2].tolist() not in net_MST[:, :2].tolist()])
            network_data = np.array([i for i in network_data if i[:2].tolist() in net_MST[:, :2].tolist()])
        else:
            alt_data = np.array([])

        self.logger.info('[Y-LineageTracker] Calculating median haplotypes finished')

        return network_data, alt_data, sp_sequence_df

    # prune network for visualization
    def _get_network_for_plot(self, network_data, alt_data, hap_num, alt_type):

        # specify network data according to type of alternative link
        if alt_data.size != 0:
            alt_type = alt_type
            pruned_network_data = network_data.copy()
            max_number = np.max(pruned_network_data[:, :2]) + 1
            while True:
                to_be_deleted_row = []
                for i in range(hap_num, max_number):
                    if np.count_nonzero(pruned_network_data[:, :2]==i) == 1:
                        row_index = np.where(pruned_network_data[:, :2]==i)[0]
                        to_be_deleted_row.append(row_index)
                pruned_network_data = np.delete(pruned_network_data, to_be_deleted_row, axis=0)
                if len(to_be_deleted_row) == 0:
                    break
            # get network data for visulization
            if alt_type == 'hide':
                network_data = pruned_network_data
        else:
            alt_type = 'hide'
            pruned_network_data = network_data.copy()

        self.net_info.write('#HapLinks:\n')
        [self.net_info.write('\t'.join(i)+'\n') for i in pruned_network_data.astype(str)]
        self.net_info.close()
        node_num = len(np.unique(network_data[:, :2]))

        return network_data, node_num

    # get mutation number
    def _get_mutations(self, network_line, sp_sequence_df, STR_code_pos):

        seq1 = sp_sequence_df.at[network_line[0], 'Haplotype']
        seq2 = sp_sequence_df.at[network_line[1], 'Haplotype']
        mutation = [STR_code_pos[n] for n, s in enumerate(zip(seq1, seq2)) if s[0] != s[1]]

        return mutation

    # write and output fdi file
    def _get_fdi_file(self, network_data, torso_num, node_num, sequence_df, sp_sequence_df, sample_sp_code, pos, hap_size, hap_freq, STR_code_pos, pops):

        fdi_file = open(self.path+'.fdi', 'w')
        medians = sorted(list(set(set(range(len(sample_sp_code), sp_sequence_df.index.size)) & set(pos.keys()))))

        # head info
        center_x,  center_y = np.mean(np.array(list(pos.values())).T, axis=1)
        fdi_file.write('SHOW_NODENAMES;TRUE\n'\
                       'NODES_PROPORTIONAL;TRUE\n'\
                       'SHOW_ONLY_TORSO;FALSE\n'\
                       'SHOW_MEDIANS;FALSE\n'\
                       'SHOW_MEDIAN_NAMES;TRUE\n'\
                       'SHOW_CHARACTERS;FALSE\n'\
                       'SHOW_LINES;FALSE\n'\
                       'MIN_CIRC_RADIUS;4\n'\
                       'MAX_CIRC_RADIUS;50\n'\
                       'MED_RADIUS;3\n'\
                       '\n'\
                       'SPECIAL_FILE_TYPE;FALSE\n'\
                       'DIAGRAM_CENTER;%d;%d\n'\
                       'IMAGE_CENTER;%d;%d\n'\
                       % (center_x, center_y, center_x, center_y))

        # taxon info
        from FilesIO import set_color_num
        from matplotlib import colors
        rgb_colors = [list(map(lambda x: round(x*255), c)) for c in set_color_num(len(pops))]
        colors = [i[0]+i[1]*256+i[2]*256*256 for i in rgb_colors]
        fdi_file.write('NUMBER_OF_TAXA;%d\n' % (sequence_df.index.size+len(medians))) # including median nodes
        plot_num = 0
        for name in sequence_df.index:
            if name in sample_sp_code.keys():
                size = hap_size[plot_num]
                freq = hap_freq[plot_num]
                pos_x, pos_y = list(map(round, pos[plot_num]))
                active = 'TRUE'
                plot_num += 1
            else:
                size = 1
                pos_x, pos_y = 0, 0
                active = 'FALSE'
                freq = None
            taxon_info = 'TAXON_NAME;%s;'\
                         'TAXON_NAME_HEIGHT;6;'\
                         'TAXON_NAME_COLOR;0;'\
                         'TAXON_FREQUENCY;%d;'\
                         'TAXON_ORIG_FREQUENCY;1;'\
                         'TAXON_GEOGRAPHY;;'\
                         'TAXON_PHENOTYPE;;'\
                         'TAXON_LINEAGE;;'\
                         'TAXON_GROUP1;Group1;;'\
                         'TAXON_GROUP2;Group2;;'\
                         'TAXON_GROUP3;Group3;;'\
                         'TAXON_X;%d;'\
                         'TAXON_Y;%d;'\
                         % (name, size, pos_x, pos_y)
            if freq:
                pie_info = ''
                colored_pie_num = 0
                for pie_num, pie_freq in enumerate(freq):
                    if pie_freq == 0:
                        continue
                    else:
                        colored_pie_num += 1
                        pie_color_code = colors[pie_num]
                        pie_info += 'TAXON_COLOR_PIE%d;%d;'\
                                    'TAXON_PIE_FREQUENCY%d;%d;'\
                                    'TAXON_STYLE_PIE%d;SOLID;'\
                                    % (colored_pie_num, pie_color_code, colored_pie_num, pie_freq, colored_pie_num)
            else:
                pie_info = 'TAXON_COLOR_PIE1;65535;'\
                           'TAXON_PIE_FREQUENCY1;1;'\
                           'TAXON_STYLE_PIE1;SOLID;'\

            line_info = 'TAXON_LINE_WIDTH;1;'\
                        'TAXON_LINE_COLOR;0;'\
                        'TAXON_LINE_STYLE;SOLID;'\
                        'TAXON_ACTIVE;%s\n'\
                        % (active)

            fdi_file.write(taxon_info + pie_info + line_info)

        # median info
        for n, m in enumerate(medians):
            pos_x, pos_y = list(map(round, pos[m]))
            median_info = 'TAXON_NAME;mv%d;'\
                          'TAXON_NAME_HEIGHT;6;'\
                          'TAXON_NAME_COLOR;0;'\
                          'TAXON_FREQUENCY;0;'\
                          'TAXON_ORIG_FREQUENCY;0;'\
                          'TAXON_GEOGRAPHY;;'\
                          'TAXON_PHENOTYPE;;'\
                          'TAXON_LINEAGE;;'\
                          'TAXON_GROUP1;Group1;;'\
                          'TAXON_GROUP2;Group2;;'\
                          'TAXON_GROUP3;Group3;;'\
                          'TAXON_X;%d;'\
                          'TAXON_Y;%d;'\
                          'TAXON_COLOR_PIE1;255;'\
                          'TAXON_PIE_FREQUENCY1;0;'\
                          'TAXON_STYLE_PIE1;SOLID;'\
                          'TAXON_LINE_WIDTH;1;'\
                          'TAXON_LINE_COLOR;0;'\
                          'TAXON_LINE_STYLE;'\
                          'SOLID;TAXON_ACTIVE;TRUE\n'\
                          % (n, pos_x, pos_y)
            fdi_file.write(median_info)

        fdi_file.write('\n')
        fdi_file.write('NUMBER_OF_LINKS;%d\n' % len(network_data))
        for n, link in enumerate(network_data):
            if n < torso_num:
                in_torso = 'TRUE'
            else:
                in_torso = 'FALSE'
            link_names = [list(sample_sp_code.keys())[name_code] if name_code < len(sample_sp_code) else 'mv'+str(medians.index(name_code)) for name_code in link[:2]]
            link_info = 'LINK_TAXON1;%s;'\
                        'LINK_TAXON2;%s;'\
                        'LINK_IN_TORSO;%s;'\
                        'LINK_TORSO_PART;1;'\
                        'NUMBER_OF_MUTATIONS;%d;'\
                        % (link_names[0], link_names[1], in_torso ,link[2])
            mutation_data = self._get_mutations(link, sp_sequence_df, STR_code_pos)
            mutation_info = ''
            for mutation_n, mutation in enumerate(mutation_data):
                mutation_info += 'MUTATION%d_NAME;%s;'\
                                 'MUTATION%d_NAME_HEIGHT;6;'\
                                 'MUTATION%d_NAME_COLOR;255;'\
                                 'LINK_LINE_WIDTH;1;'\
                                 'LINK_LINE_STYLE;SOLID;'\
                                 'LINK_LINE_COLOR;8421504;'\
                                 'LINK_ACTIVE;TRUE'\
                                 % (mutation_n+1, mutation, mutation_n+1, mutation_n+1)
            fdi_file.write(link_info+mutation_info+'\n')

        # other info
        # median nodes not included
        fdi_file.write('\n')
        other_info = 'NUMBER_OF_HELP_POINT_LINKS;0\n'\
                     '\n'\
                     'LONG_EDGES;0\n'\
                     'TOTAL_TAXA;%d\n'\
                     'TOTAL_HAPLOTYPES;%d\n'\
                     % (sequence_df.index.size, len(sample_sp_code))
        fdi_file.write(other_info)

        # equal info
        equal_info = ''
        equal_num = 0
        for sample, hap in sample_sp_code.items():
            equal_samples = sequence_df[sequence_df['Haplotype']==hap].index.tolist()
            for equal in equal_samples:
                if equal != sample:
                    equal_num += 1
                    equal_info += 'EQUIVALENT_TAXA;%s;%s\n' % (sample, equal)
        fdi_file.write('NUMBER_EQUIVALENT_TAXA;%d\n' % equal_num)
        fdi_file.write(equal_info)

        # end info
        end_info = 'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Phenotype;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Geography;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Lineage;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Group1;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Group2;\n'\
                   'COLOR_SCHEME;SCHEME_ACTIVE;FALSE;SCHEME_TYPE;Group3;\n'
        fdi_file.write(end_info)

        fdi_file.close()

    # summarize parameters for figure
    def _get_figure_parameters(self, sequence_df, population_data, network_data, alt_data, node_num, original_network_size, alt_type, euqal_size, label):

        haplotypes = sorted(list(set(sequence_df['Haplotype'])))
        # calculate frequency of each haplotypes
        if euqal_size:
            hap_size = [1]*len(haplotypes)
        else:
            hap_size = [sequence_df['Haplotype'].tolist().count(hap) for hap in haplotypes]
            hap_size_per = [size / sequence_df.index.size for size in hap_size ]

        if isinstance(population_data, pd.DataFrame):
            pops = sorted(list(set(population_data[label])))
            # calculate haplotype frequency in each population
            hap_pop_data = pd.DataFrame(index=range(len(haplotypes)), columns=pops)
            for pop in pops:
                pop_inds = population_data[population_data[label]==pop]['SampleID'].tolist()
                pop_data = sequence_df.loc[pop_inds]
                hap_pop_data[pop] = [pop_data['Haplotype'].tolist().count(hap) for hap in haplotypes]
            hap_freq = [hap_pop_data.loc[i].tolist() for i in hap_pop_data.index]
            hap_freq_per = [(hap_pop_data.loc[i]/sum(hap_pop_data.loc[i])).round(6).tolist() for i in hap_pop_data.index]
        else:
            hap_freq_per = None
            pops = None

        # set length of each node
        length_data = np.empty(shape=(0, 3))
        num = np.max(network_data[:, :2])+1
        to_be_deleted_row = []
        for row_index, i in enumerate(network_data):
            step = i[2]
            if step > 1:
                start = i[0]
                end = i[1]
                for i in range(step):
                    if i == 0:
                        length = [start, num, 1]
                    elif i == step-1:
                        length = [num, end, 1]
                    else:
                        length = [num, num+1, 1]
                        num += 1
                    length_data = np.concatenate((length_data, np.array([length])))
                num += 1
                to_be_deleted_row.append(row_index)
                original_network_size -= 1
        network_data = np.delete(network_data, to_be_deleted_row, axis=0)
        length_data = length_data.astype(int)

        # set figure size and network class
        import warnings
        warnings.filterwarnings('ignore')
        font_files = font_manager.findSystemFonts(fontpaths=[os.path.split(os.path.realpath(__file__))[0] + '/sans-serif'])
        font_list = font_manager.createFontList(font_files)
        font_manager.fontManager.ttflist.extend(font_list)
        plt.rcParams['font.family']= 'sans-serif'
        plt.rcParams['font.sans-serif'] = ['Arial'] # font
        plt.rcParams['legend.handlelength'] = 1
        plt.rcParams['legend.handleheight'] = 1
        fig_len = len(length_data)/3
        fig, ax = plt.subplots(figsize=(fig_len, fig_len))
        trans = ax.transData.transform
        trans2 = fig.transFigure.inverted().transform
        DG = nx.DiGraph()

        # construct network by adding nodes and edges
        style_list = []
        for i, j in enumerate(network_data):
            if i < original_network_size:
                DG.add_edge(j[0], j[1])
                style_list.append('solid')
            else:
                DG.add_edge(j[0], j[1])
                if alt_type != 'dashed':
                    style_list.append('dashed')
                else:
                    style_list.append('solid')
        # construct network by adding node of length
        for i in length_data:
            DG.add_edge(i[0], i[1])
            style_list.append('solid')

        try:
            import pygraphviz as pg
            pos = nx.nx_agraph.graphviz_layout(DG, prog='neato')
        except ImportError:
            pre_pos = nx.spring_layout(DG)
            pos = {}
            for i in pre_pos.keys():
                pos[i] = trans(pre_pos[i])

        return length_data, DG, pos, hap_size, hap_freq, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len

    # plot a tree and replace nodes to pies
    def plot_network(self, length_data, DG, pos, node_num, median_size_rate, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len, label):

        import warnings
        warnings.filterwarnings('ignore')

        # set figure parameters
        # mdian haplotype size
        base_size = 5*20
        median_size = base_size * median_size_rate
        # all node size
        hap_num = len(hap_size_per)
        max_size = max(hap_size_per)
        size_rate = base_size/max_size
        size = [hap_size_per[i]*size_rate if i < hap_num else median_size for i in DG.nodes][:node_num]
        size = size + [0.2] * length_data.size

        try:
            # color and line type
            if hap_freq_per:
                from matplotlib.patches import Patch
                from FilesIO import set_color_num
                nx.draw(DG, pos=pos, arrows=False, node_color='black', node_size=min(size), style=style_list)
                # set colors
                colors = set_color_num(len(pops))
                legend_elements = []
                for color, pop in zip(colors , pops):
                    legend_elements.append(Patch(facecolor=color, edgecolor='black', label=pop))
                plt.legend(handles=legend_elements, frameon=False, bbox_to_anchor=(-0.04, 0.5), prop={'size': fig_len})
                pie_size = 0.025
                move = pie_size/2.0
                for i in DG:
                    if i < hap_num:
                        xx, yy = trans(pos[i])
                        xa, ya = trans2((xx,yy))
                        pie_ax = plt.axes([xa-move, ya-move, pie_size, pie_size])
                        pie_ax.set_aspect('equal')
                        fracs = hap_freq_per[i]
                        if 1 in fracs:
                            one_colors= [colors[fracs.index(1)]]
                            one_fracs = [1]
                            pie_ax.pie(one_fracs, colors=one_colors, radius=(1/(max_size/hap_size_per[i]))**0.5, wedgeprops={'edgecolor':'black','linewidth': 0.5, 'linestyle': 'solid'})
                        else:
                            pie_ax.pie(fracs, colors=colors, radius=(1/(max_size/hap_size_per[i]))**0.5, wedgeprops={'edgecolor':'black','linewidth': 0.5, 'linestyle': 'solid'})
            else:
                colors = ['yellow' if i < hap_num else 'black' for i in DG.nodes] # color
                nx.draw(DG, pos=pos, arrows=False, node_color=colors, node_size=size, style=style_list, linewidths=1)
                node_ax = plt.gca() # to get the current axis
                node_ax.collections[0].set_edgecolor('black')
                node_ax.collections[0].set_linewidth(0.5)

            # output
            if label:
                plt.savefig(self.path+'.net.fig.%s.pdf' % label, dpi=100, bbox_inches='tight')
            else:
                plt.savefig(self.path+'.net.fig.pdf', dpi=100, bbox_inches='tight')
            plt.close()
        except AttributeError:
            print('[Y-LineageTracker] Cannot generate figure since matplotlib.cbook version bugs')

    # main function of network analysis
    def calculate_and_output(self, haplotype_data, population_data, arguments, label):

        # convert data to 0-1 code and calculate distance matrix
        sequence_df, STR_code_pos = self._convert_allel_to_binary(haplotype_data, arguments.filter)
        distance_matrix, sp_sequence_df, sample_sp_code = self._calculate_distance_matrix(sequence_df)
        # calculate network
        self.logger.info('[Y-LineageTracker] Calculating median haplotypes')
        network_data, alt_data, sp_sequence_df2 = self._calculate_network(distance_matrix, sp_sequence_df, arguments.treetype, arguments.search)
        plot_network_data, node_num = self._get_network_for_plot(network_data, alt_data, sp_sequence_df.index.size, arguments.alt)

        # visulization
        self.logger.info('[Y-LineageTracker] Start to plot network')

        if label != 'Group':
            length_data, DG, pos, hap_size, hap_freq, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len = self._get_figure_parameters(sequence_df, population_data, plot_network_data, alt_data, node_num, network_data.size, arguments.alt, arguments.euqalsize, label)
            if arguments.fdi:
                self._get_fdi_file(plot_network_data, plot_network_data.size, node_num, sequence_df, sp_sequence_df2, sample_sp_code, pos, hap_size, hap_freq, STR_code_pos, pops)
            self.plot_network(length_data, DG, pos, node_num, arguments.mediansize, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len, label)
        else:
            label = ['Population', 'Group']
            for i in label:
                length_data, DG, pos, hap_size, hap_freq, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len = self._get_figure_parameters(sequence_df, population_data, plot_network_data, alt_data, node_num, network_data.size, arguments.alt, arguments.euqalsize, label)
                if arguments.fdi:
                    self._get_fdi_file(plot_network_data, plot_network_data.size, node_num, sequence_df, sp_sequence_df2, sample_sp_code, pos, hap_size, hap_freq, STR_code_pos, pops)
                self.plot_network(length_data, DG, pos, node_num, arguments.mediansize, hap_size_per, hap_freq_per, style_list, pops, trans, trans2, fig_len, label)

        self.logger.info('[Y-LineageTracker] Network analysis finished')

def main():

    start = time.perf_counter()
    arguments = net_parser()

    from FilesIO import check_seq_matrix_input, check_hap_input, set_haplotype_index, check_population, check_overlap, check_pop_header, time_count
    # check input files
    input, format = check_seq_matrix_input(arguments.seq, arguments.format, arguments.matrix)

    # check haplotype data
    haplotype_data = check_hap_input(input, 'haplotype', format)

    # check filter argument
    check_misisng_rate(arguments.filter)

    from FilesIO import get_out_path
    # get path of output files
    path = get_out_path(arguments.output, 'Net')
    log_file = path + '.NetLog.log'

    # set log file
    set_log(arguments, log_file)

    # check if there is a population file
    if arguments.population:
        population_data = check_population(arguments.population)
        population_data = check_pop_header(population_data)
        haplotype_data, population_data = check_overlap(haplotype_data, population_data)
        label = 'Population'
        if 'Group' in population_data.columns:
            label = 'Group'
    else:
        population_data = None
        label = ''

    # set index of haplotype data
    haplotype_data = set_haplotype_index(haplotype_data, format)

    # calculate network and visualize result
    net = HaploNetwork(arguments.gap, path)
    net.calculate_and_output(haplotype_data, population_data, arguments, label)

    time_count(start)

if __name__ == '__main__':
    main()

{'set_color_num': 'FilesIO.set_color_num', 'colors': 'matplotlib.colors', 'warnings': 'warnings', 'pg': 'pygraphviz', 'Patch': 'matplotlib.patches.Patch'}