import logging  # Logging functionality
from os import truncate

import boto3  # import the boto api
import subprocess
import time  # pausing the interface
from tinydb import TinyDB, Query, where
#import s3_check_create  # seperate python file to check for s3 bucket
import sys  # this allows you to use the sys.exit command to quit/logout of the application
from datetime import datetime, timedelta

# Menu for new webserver
db = TinyDB('db.json')
# Global declaration of ec2 and cloudwatch
cloudwatch = boto3.resource('cloudwatch')
ec2 = boto3.resource('ec2')
ec2Client = boto3.client('ec2')
filename='log/aws_assignment.log'
# Configuration for the logfile https://docs.python.org/3/howto/logging.html + https://www.pylenin.com/blogs/python-logging-guide/
logger_format= "%(asctime)s::%(levelname)s::%(name)s::"\
             "%(filename)s::%(lineno)d::%(message)s"
logging.basicConfig(
    filename= filename,
    level=logging.DEBUG,
    format = logger_format,
    datefmt='%d/%m/%Y %I:%M:%S %p'
)
def main():    #Main function to call the main menu
    logging.info('Program started')
    menu()
def menu():
    logging.info('Main menu selected')
    print("\n\n\n              ************MAIN MENU**************")

    time.sleep(0.02)
    print()

    choice = input("""
                      A: Instance Menu
                      B: Monitoring menu
                      C: Open Logfile for information
                      D: VPC menu
                      -------------------
                      Q: Quit/Log Out

                      Please enter your choice: """)

    if choice == "A" or choice == "a":
        instance_menu()
    elif choice == "B" or choice == "b":
        monitor_menu()
    elif choice == "C" or choice == "c":
        open_logfile()
    elif choice == "D" or choice == "d":
            vpc_menu()

    elif choice == "Q" or choice == "q":
        logging.info("Exiting of program")
        sys.exit()
    else:
        print("You must only select either A,B,C, or D.")
        print("Please try again")
        menu()

def createNewInstance():
    ec2 = boto3.resource('ec2')
    print("\nStarting instance creation process, please be patient")
    try:
        TAG_SPEC = [
        {
        "ResourceType":"instance",
        "Tags": [
                {
                    "Key": "Name",
                    "Value": "BastionServer"
                }
            ]
        }
        ]
        NET_SPEC = [
            {
                "AssociatePublicIpAddress": "True"
            }
        ]
        instance = ec2.create_instances(
            ImageId='ami-099a8245f5daa82bf',  # Default instance
            MinCount=1,
            MaxCount=1,
            InstanceType='t2.nano',  # t2 nano default or micro depending upon selection
            KeyName='kp20201',
            TagSpecifications = TAG_SPEC,
            SecurityGroupIds=['sg-0f20389f889bed0a3'],  # one of my security groups that has http and ssh
            SubnetId = 'subnet-06a71d27240482424', #Launch into the public subnet
            #NetworkInterfaces = NET_SPEC
        )
        print("Instance ID:" + instance[0].id + " being created. Please be patient!")
    except:
        logging.error("Couldn't create an instance")
        print("Couldn't create an instance")
        input("\nPress Enter to continue...")
        instance_menu()
    # instance created but no tags inserted.
    try:
        instance[0].wait_until_running()
        print("Instance running")
    except:
        print("Cannot check for instanance running")
        logging.warning("Unable to check for instance running")
    try:
        print("Instance running: " + instance[0].id)
    except:
        logging.warning("No instance running on the system due to errors")
        print("No instance running")
    input("\nPress Enter to continue...")
    instance_menu()
def createNewInstanceDbServer():
    ec2 = boto3.resource('ec2')
    print("\nStarting instance creation process, please be patient")
    try:
        TAG_SPEC = [
        {
        "ResourceType":"instance",
        "Tags": [
                {
                    "Key": "Name",
                    "Value": "dbServer"
                }
            ]
        }
        ]
        NET_SPEC = [
            {
                "AssociatePublicIpAddress": "True"
            }
        ]
        instance = ec2.create_instances(
            ImageId='ami-000f250cb403819f1',  # Default instance
            MinCount=1,
            MaxCount=1,
            InstanceType='t2.nano',  # t2 nano default or micro depending upon selection
            KeyName='kp20201',
            TagSpecifications = TAG_SPEC,
            SecurityGroupIds=['sg-0d215f7e40740063a'],  # one of my security groups that has http and ssh
            SubnetId = 'subnet-0113caae4b929264e', #Launch into the private subnet
            PrivateIpAddress = '20.0.4.86',
            UserData = '#!/bin/bash su \'mongod -dbpath /home/ec2-user/db --bind_ip_all \''
         )
        print("Instance ID:" + instance[0].id + " being created. Please be patient!")
    except Exception as e:
        print(e)
        logging.error("Couldn't create an instance")
        print("Couldn't create an instance")
        input("\nPress Enter to continue...")
        instance_menu()
    # instance created but no tags inserted.
    try:
        instance[0].wait_until_running()
        print("Instance running")
    except:
        print("Cannot check for instanance running")
        logging.warning("Unable to check for instance running")
    try:
        print("Instance running: " + instance[0].id)
    except:
        logging.warning("No instance running on the system due to errors")
        print("No instance running")

    input("\nPress Enter to continue...")
    instance_menu()
def list_all_instance():  # Listing of all instances in all statuses
    ec2 = boto3.resource('ec2')
    print("\nAttempting to list instances, please be patient")
    try:
        instance_list = []
        for instance in ec2.instances.all():
            print(instance.id, instance.state, instance.public_ip_address, instance.tags)
            instance_list.append(instance)
        logging.warning("Instance list created without issue")
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error handling instances more than likely due to connection")
    input("\nPress Enter to continue...")
    instance_menu()
def instance_menu():
    logging.info('Instance Menu Selected')
    print("\n\n\n              ************INSTANCE MENU**************")
    time.sleep(0.02)
    print()

    choice = input("""
                      A: Create a new Bastion Instance (in development)
                      B: List all instances (any state)
                      C: Create a new dbServer instance (in development
                    -----------------------------------------------
                      D: Terminate instance
                    -----------------------------------------------
                    -----------------------------------------------
                      Q: Back to Main Menu

                      Please enter your choice: """)

    if choice == "A" or choice == "a":
        logging.info('Create new instance selected')
        instance_menu()
        #createNewInstance()
    elif choice == "B" or choice == "b":
        list_all_instance()   #list all instances - different function to show all statuses
    elif choice == "C" or choice == "c":
        instance_menu()
        #createNewInstanceDbServer()   #create a dbServer instance
    elif choice == "D" or choice == "d":
        quitInstance()
    elif choice == "E" or choice == "e":
        instance_menu()
    elif choice == "Q" or choice == "q":
        menu()
    else:
        print("You must only select either A,B,C, or D.")
        print("Please try again")
        instance_menu()

def vpc_menu():
    logging.info('VPC Menu Selected')
    print("\n\n\n              ************Automated Systems**************")
    time.sleep(0.02)
    print()

    choice = input("""
                      A: Create an automated system *Assignment 2 setup
                      B:
                      C:
                    -----------------------------------------------
                      D: Delete and automated setups
                    -----------------------------------------------
                      E: List all vpc's
                    -----------------------------------------------
                      Q: Back to Main Menu

                      Please enter your choice:
                      """)
    if choice == "A" or choice == "a":
        logging.info('Create new VPC selected')
        new_vpc()
    elif choice == "D" or choice == "d":
        logging.info("Delete automated setup")
        deleteautomated()
    elif choice == "E" or choice == "e":
        listvpcs()
    elif choice == "Q" or choice == "q":
        menu()
    else:
        print("You must only select either A,B,C, or D.")
        print("Please try again")
        vpc_menu()
def listvpcs():

    fltr1 = [{'Name': 'tag:Name','Values':['Assignment 2 Boto3']}]
    #vpclist=list(ec2.vpcs.filter(Filters=fltr))
    vpclist = list(ec2.vpcs.filter())
    #print(vpclist[0].id)
    print(vpclist)
    for i in vpclist:
        vpc_list = ec2Client.describe_vpcs(
            VpcIds=[
                i.id,
            ],
        )
        print(vpc_list)
    choice = input("\nPress Enter to continue...")
    vpc_menu()
def deleteautomated():
    print("This needs to be created, note that it will go through any setups in the db and then check for live systems \n and delete as required")
    #Todo: create deleteautomated script by checking for a live system
def new_vpc():
    print("This creates an automated setup")
    ec2 = boto3.resource('ec2')
    vpc = ec2.create_vpc(CidrBlock='20.0.0.0/16')
    print("VPC has been created: ")
    print(vpc)
    vpc.create_tags(Tags=[{"Key": "Name", "Value": "Assignment 2 Boto3"}])
    vpc.wait_until_available()

    ec2Client.modify_vpc_attribute( VpcId = vpc.id , EnableDnsSupport = { 'Value': True } )
    ec2Client.modify_vpc_attribute( VpcId = vpc.id , EnableDnsHostnames = { 'Value': True } )
    # create an internet gateway and attach it to VPC
    internetgateway = ec2.create_internet_gateway()
    print("Internet gateway setup: ")
    print(internetgateway)
    #db.insert({'VPC_ID': vpc.id,'IGW_ID':internetgateway.id})
    #print(db.all())
    igw_tag = internetgateway.create_tags(Tags=[{"Key": "Name", "Value": "Assignment 2 Boto3"}])
    vpc.attach_internet_gateway(InternetGatewayId=internetgateway.id)
    print("Internet gateway attached to the vpc")
    # create a route table and a public route
    routetable = vpc.create_route_table()
    route = routetable.create_route(DestinationCidrBlock='0.0.0.0/0', GatewayId=internetgateway.id)
    print("route table completed")
    print(route)
    route_tags = routetable.create_tags(Tags=[{"Key": "Name", "Value": "Assignment 2 Boto3"}])
    # create subnet and associate it with route table
    subnet = ec2.create_subnet(CidrBlock='20.0.1.0/24', VpcId=vpc.id, AvailabilityZone='eu-west-1a')
    subnet.create_tags(
        Tags=[
            {
                'Key': 'Name',
                'Value': 'Public Subnet 1a'
            },
        ]
    )

    routetable.associate_with_subnet(SubnetId=subnet.id)
    print("Public subnet 1a created ")
    print(subnet)
    subnet2 = ec2.create_subnet(CidrBlock='20.0.2.0/24', VpcId=vpc.id, AvailabilityZone='eu-west-1b')
    subnet2.create_tags(
            Tags=[
                {
                    'Key': 'Name',
                    'Value': 'Public Subnet 1b'
                },
            ]
    )
    print("Public subnet 1b created ")
    print(subnet2)
    routetable.associate_with_subnet(SubnetId=subnet2.id)
    subnet3 = ec2.create_subnet(CidrBlock='20.0.3.0/24', VpcId=vpc.id, AvailabilityZone='eu-west-1c')
    subnet3.create_tags(
            Tags=[
                {
                    'Key': 'Name',
                    'Value': 'Public Subnet 1c'
                },
            ]
    )
    print("Public subnet 1c created ")
    print(subnet3)
    routetable.associate_with_subnet(SubnetId=subnet3.id)

    #db.insert({'VPC_ID': vpc.id,'IGW_ID':internetgateway.id,"route_table_public":routetable.id, "subnet1_id":subnet.id, "subnet2_id":subnet2.id, "subnet3_id":subnet3.id})
    #print(db.all())
    #time.sleep(15)

    # Create a security group and allow SSH inbound rule through the VPC
    securitygroup1 = ec2.create_security_group(GroupName='webserver', Description='webserverSG', VpcId=vpc.id)
    securitygroup1.authorize_ingress(CidrIp='81.24.248.5/32', IpProtocol='tcp', FromPort=22, ToPort=22)
    print("Security Group webserver created, ID =")
    print(securitygroup1.id)
    # Create a security group and allow SSH inbound rule through the VPC
    securitygroup2 = ec2.create_security_group(GroupName='bastion', Description='bastionSG', VpcId=vpc.id)
    securitygroup2.authorize_ingress(CidrIp='81.24.248.5/32', IpProtocol='tcp', FromPort=22, ToPort=22)
    print("Security Group bastion created, ID =")
    print(securitygroup2.id)
    # Create a security group and allow SSH from bastion to dbServer
    securitygroup3 = ec2.create_security_group(GroupName='dbserver', Description='dbServerSG', VpcId=vpc.id)
    print("Security Group dbserver created id:")
    print(securitygroup3)

    #Assignment of ports for communication to the dbServer
    ec2Client.authorize_security_group_ingress(GroupId=securitygroup3.id,
                                                   IpPermissions=[
                                                       {
                                                           'FromPort': 22,
                                                           'IpProtocol': 'tcp',
                                                           'ToPort': 22,
                                                           'UserIdGroupPairs': [
                                                               {
                                                                   'Description': 'SSH access from the bastion server',
                                                                   'GroupId': securitygroup2.id,
                                                               },
                                                           ],
                                                       },
                                                   ],
                                               )
    #Assignment of ports for communication to the dbServer
    ec2Client.authorize_security_group_ingress(GroupId=securitygroup3.id,
                                                       IpPermissions=[
                                                           {
                                                               'FromPort': 27017,
                                                               'IpProtocol': 'tcp',
                                                               'ToPort': 27017,
                                                               'UserIdGroupPairs': [
                                                                   {
                                                                       'Description': 'MongoDB Access from webserver',
                                                                       'GroupId': securitygroup1.id,
                                                                   },
                                                               ],
                                                           },
                                                       ],
                                                   )
    #create rule allowing for incoming traffic to the webserver from the loadbalancer
    securitygroup_lb = ec2.create_security_group(GroupName='loadbalancer', Description='loadbalancerSG', VpcId=vpc.id)
    securitygroup_lb.authorize_ingress(CidrIp='0.0.0.0/0', IpProtocol='tcp', FromPort=80, ToPort=80)
    securitygroup_lb.authorize_ingress(CidrIp='0.0.0.0/0', IpProtocol='tcp', FromPort=443, ToPort=443)
    print("Security Group loadbalancer created, ID = ")
    print(securitygroup_lb.id)
    #create rule allowing for incoming traffic to the webserver from the loadbalancer
    ec2Client.authorize_security_group_ingress(GroupId=securitygroup1.id,
                                                       IpPermissions=[
                                                           {
                                                               'FromPort': 3000,
                                                               'IpProtocol': 'tcp',
                                                               'ToPort': 3000,
                                                               'UserIdGroupPairs': [
                                                                   {
                                                                       'Description': 'Port 3000 Traffic from the loadbalancer',
                                                                       'GroupId': securitygroup_lb.id,
                                                                   },
                                                               ],
                                                           },
                                                       ],
                                                   )

    #Allocate an Elastic IP address
    eip = ec2Client.allocate_address(
    Domain='vpc',
    )
    time.sleep(2)
    eipAllocationId=eip['AllocationId']
    print("Elastic IP created: ")
    print(eip)
    time.sleep(5)
    #Create a NAT gateway
    nat = ec2Client.create_nat_gateway(
        AllocationId=eipAllocationId,
        SubnetId= subnet.id,
    )
    print("The details of the nat are the following: ")
    print(nat)
    waiter = ec2Client.get_waiter('nat_gateway_available')
    #Todo check for this waiter functionality
    natGatewayId = nat['NatGateway']['NatGatewayId']
    print("The NatGatewayId is")
    print(nat)
    print("Please wait while the NAT is being created")
    time.sleep(20)
    # create a route table and a private route
    private_routetable = vpc.create_route_table()
    route = private_routetable.create_route(DestinationCidrBlock='0.0.0.0/0', NatGatewayId=natGatewayId)
    print("The route has been created for the private subnet")
    print(route)
    route_tags = private_routetable.create_tags(Tags=[{"Key": "Name", "Value": "NAT Assignment 2 Boto3"}])
    subnet_private = ec2.create_subnet(CidrBlock='20.0.4.0/24', VpcId=vpc.id, AvailabilityZone='eu-west-1a')
    time.sleep(1)
    subnet_private.create_tags(
                Tags=[
                    {
                        'Key': 'Name',
                        'Value': 'Private Subnet 1a'
                    },
                ]
        )
    print("Private subnet created")
    print(subnet_private)
    private_routetable.associate_with_subnet(SubnetId=subnet_private.id)
    # adding the creation of the db instance to the automated script
    time.sleep(10)
    ec2 = boto3.resource('ec2')
    print("\nStarting instance creation process, please be patient")
    try:
        TAG_SPEC = [
            {
                "ResourceType": "instance",
                "Tags": [
                    {
                        "Key": "Name",
                        "Value": "dbServer"
                    }
                ]
            }
        ]
        NET_SPEC = [
            {
                "AssociatePublicIpAddress": "True"
            }
        ]
        userdata = """#!/bin/bash
                   nohup mongod -dbpath /home/ec2-user/db --bind_ip_all
        """
        instance = ec2.create_instances(
            ImageId='ami-00478fab923418f90',  # Db instance
            MinCount=1,
            MaxCount=1,
            InstanceType='t2.nano',  # t2 nano default or micro depending upon selection
            KeyName='kp20201',
            TagSpecifications=TAG_SPEC,
            SecurityGroupIds=[securitygroup3.id,],  # one of my security groups that has http and ssh
            SubnetId=subnet_private.id,  # Launch into the private subnet
            PrivateIpAddress='20.0.4.86',
           # UserData=userdata
           # Note that mongod is setup to auto run on this AMI instance [sudo systemctl enable mongod]
        )
        print("Instance ID:" + instance[0].id + " being created. Please be patient!")
    except Exception as e:
        print(e)
        logging.error("Couldn't create an instance")
        print("Couldn't create an instance")
        input("\nPress Enter to continue...")
        instance_menu()
    #https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/elbv2.html
    elb = boto3.client('elbv2')
    #create the target group
    tg = elb.create_target_group(
        Name='Assignment2-boto3-tg',
        Port=3000,
        Protocol='HTTP',
        VpcId=vpc.id,
    )
    print("Rarget group details: ")
    print(tg)
    targetgroup_arn= tg['TargetGroups'][0]['TargetGroupArn']
    #Create the loadbalancer
    elb_info = elb.create_load_balancer(
        Name='Boto3-load-balancer',
        Scheme='internet-facing',
        Subnets=[
            subnet.id,
            subnet2.id,
            subnet3.id,
        ],
        SecurityGroups= [
                        securitygroup_lb.id,
        ],
        Type= 'application',
    )
    print("Load balancer information: ")
    print(elb)
    elb_arn= elb_info['LoadBalancers'][0]['LoadBalancerArn']
    #Create the listner on the LB
    elb_listner1 = elb.create_listener(
        DefaultActions=[
            {
                'TargetGroupArn': targetgroup_arn,
                'Type': 'forward',
            },
        ],
        LoadBalancerArn=elb_arn,
        Port=80,
        Protocol='HTTP',
    )
    print("The listner for the elb is the following: ")
    print(elb_listner1)
    elb_listner2 = elb.create_listener(
        Certificates=[
            {
                'CertificateArn': 'arn:aws:acm:eu-west-1:013355473762:certificate/6da9509f-05ad-47a1-96c9-98dc9bf8a4a4',
            },
        ],
        DefaultActions=[
            {
                'TargetGroupArn': targetgroup_arn,
                'Type': 'forward',
            },
        ],
        LoadBalancerArn=elb_arn,
        Port=443,
        Protocol='HTTPS',
        SslPolicy='ELBSecurityPolicy-2016-08',
    )
    print("The listner for the elb is the following: ")
    print(elb_listner1)

    db.insert(
                     {
                     'VPC_ID': vpc.id,
                     'IGW_ID':internetgateway.id,
                     'route_table_public':routetable.id,
                     'subnet1_id':subnet.id, 'subnet2_id':subnet2.id, 'subnet3_id':subnet3.id,
                     'securitygroup1':securitygroup1.id, 'securitygroup2':securitygroup2.id, 'securitygroup3':securitygroup3.id, 'securitygrouplb':securitygroup_lb.id, 'subnetprivate_id': subnet_private.id,
                     'EipAllocation_id': eipAllocationId,
                     'NatGatewayId': natGatewayId,
                     'private_routetable': private_routetable.id,
                     'targetgroup_arn': targetgroup_arn,
                     'elb_arn':elb_arn,
                     }
                   )
    test = db.all()
    #print("VPC_ID" +test[0]['VPC_ID'])

    #low level client representing Auto Scaling
    ec2auto = boto3.client('autoscaling')
    userdatalc = """#!/bin/bash
                  su -ec2user -c 'cd /home/ec2-user/donation-web-10/
                  node index.js'"""
    lc = ec2auto.create_launch_configuration(
            ImageId='ami-0e3bad3ef579c546e',
            InstanceType='t2.nano',
            LaunchConfigurationName='Boto3AssignmentLC',
            SecurityGroups=[
               securitygroup1.id,
            ],
            UserData = userdatalc,
            InstanceMonitoring={'Enabled':True},
            AssociatePublicIpAddress=True,
    )
    print("This is the LC:")
    print(lc)
    asgsubnets= subnet.id +","+subnet2.id+","+subnet3.id
    time.sleep(20)
    asg = ec2auto.create_auto_scaling_group(
            AutoScalingGroupName='Boto3AssignmentASG',
            LaunchConfigurationName='Boto3AssignmentLC',
            AvailabilityZones=[
                    'eu-west-1a','eu-west-1b','eu-west-1c'
                ],
            MaxSize=3,
            MinSize=1,
            HealthCheckType='EC2',
            VPCZoneIdentifier= asgsubnets,
            Tags=[
                    {
                        'ResourceId': 'Boto3AssignmentASG',
                        'ResourceType': 'auto-scaling-group',
                        'Key': 'Name',
                        'Value': 'Boto3AssignmentASG',
                        'PropagateAtLaunch': True,
                    },
                ],
        )
    print("asg printout: ")
    print(asg)
    response = ec2auto.attach_load_balancer_target_groups(
         AutoScalingGroupName='Boto3AssignmentASG',
         TargetGroupARNs=[
             targetgroup_arn,
         ]
    )
    print(response)

    #ami-03778245a5d9f1082

    #print(test.elb_arn)
            #print(db.all())
            #time.sleep(15)
    print("------Take notice------------")
    print("*******")
    print("Create a homemade bastion server for connection with dbserver using the bastionsg and public ip")
    print("*******")
    print("-----------------------------")
    choice = input("\nPress Enter to continue...")

    # TODO: add the asg and lc details to the db information
    # TODO: add try/except to the creation of all the items e.g. vpc etc
    # TODO: create dbserver that will run up from ami instance
    # TODO: create a new AMI for webserver instance with a static internal ip for the dbserver
    # TODO: Add details to the logger of creation of vpc...or errors

    vpc_menu()

def monitor_menu():
    print("\n\n\n              ************MONITORING MENU**************")
    time.sleep(0.02)
    print()

    choice = input("""
                      A: Monitor the CPU utilisation on an instance
                      B: Set Alarm on instance (note NetworkIn less than 30k)
                      ------------------------
                      C: Set custom monitoring on EC2 instance to cloudwatch
                      D: Get custom data back from EC2 instance
                      ------------------------
                      E: Search for all monitored instances for CPU utilisation
                      F: Search for current monitored instances for CPU utilisation
                      Q: Back to Main Menu

                      Please enter your choice: """)

    if choice == "A" or choice == "a":
        select_monitor()
    elif choice == "B" or choice == "b":
        cloudwatch_alarm()
    elif choice == "C" or choice == "c":
        pushmonitoring()
    elif choice == "D" or choice == "d":
        custom_monitoring()
    elif choice == "E" or choice == "e":
        monitor_db()
    elif choice == "F" or choice == "f":
        monitor_specific_db()
    elif choice == "Q" or choice == "q":
        menu()
    else:
        print("You must only select either A,B,C, or D.")
        print("Please try again")
        monitor_menu()
def monitor_db():
    InstanceMonitored = Query()
    print(db.search(InstanceMonitored.AverageutilisationPercentage.exists()))
    input("\nPress Enter to continue...")
    monitor_menu()
def monitor_specific_db():
    instance_list = []
    try:
        instance_list = instance_listing(
            ['running'])  # function that takes in the instance status and returns an arrany of instance id's
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error searching for instances:")
    if not instance_list:  # check for an empty array i.e. no running instances
        print("No running instances")
        logging.info("No running instances")
        monitor_menu()
    else:
        choice = input("""Please select the running instance to view CPU utilisation:""")

        try:
            selectedinstance = instance_list[int(choice)]
            InstanceMonitored = Query()
            result = db.search(InstanceMonitored.InstanceMonitored == selectedinstance.id)
            if not result:
                print("No results for this instance")
            else:
                for i in result:
                   print(i)
            input("\nPress Enter to continue...")
            monitor_menu()
        except Exception as e:
            print(e)
            logging.warning("Issue with choice entry as no data yet for the instance")

def select_monitor():
    instance_list = []
    try:
        instance_list = instance_listing(['running'])   #function that takes in the instance status and returns an arrany of instance id's
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error searching for instances:")
    if not instance_list:  # check for an empty array i.e. no running instances
        print("No running instances")
        logging.info("No running instances")
        monitor_menu()
    else:
        choice = input("""Please select the instance number to monitor:""")

        try:
            selectedinstance = instance_list[int(choice)]
            selectedinstance.monitor()

            metric_iterator = cloudwatch.metrics.filter(Namespace='AWS/EC2',
                                                        MetricName='CPUUtilization',
                                                        Dimensions=[{'Name':'InstanceId', 'Value': selectedinstance.id}])
            metric = list(metric_iterator)[0]    # extract first (only) element
            print(metric_iterator)
            response = metric.get_statistics(StartTime = datetime.utcnow() - timedelta(minutes=5),   # 5 minutes ago
                                             EndTime=datetime.utcnow(),                              # now
                                             Period=300,                                             # 5 min intervals
                                             Statistics=['Average'])
            print("Average CPU utilisation:", response['Datapoints'][0]['Average'], response['Datapoints'][0]['Unit'])
            avg = response['Datapoints'][0]['Average']
            db.insert({'InstanceMonitored': selectedinstance.id,'AverageutilisationPercentage':avg, 'Date_time':datetime.now().strftime('%Y-%m-%d %H:%M:%S')})
            input("\nPress Enter to continue...")
            if choice == "R" or choice == "r":
                select_monitor()
            else:
               monitor_menu()
        except Exception as e:
            print(e)
            logging.warning("Issue with choice entry as no data yet for the instance")
            print("Issue with choice entry as no data yet for the instance")
            choice = input("\nPress Enter to continue...or R to repeat")
            if choice == "R" or choice == "r":
                select_monitor()
            else:
               monitor_menu()
def instance_listing(status):
    i = 0
    instance_list = []
    filters = [
        {
            'Name': 'instance-state-name',
            'Values': status
        }
    ]
    logging.info("Attempting to create a list of instances")

    try:

        for instance in ec2.instances.filter(Filters=filters):
            instance_list.append(instance)
            try:
                for tag in instance.tags:  # As AWS stores tags as key and value
                    print("[%d]" % (i) + instance.id + " Tag Key: " + tag['Key'] + " Value: " + tag['Value'])
            except:
                print("[%d]" % (i) + instance.id + " No Tags")
            logging.info(instance.id + instance.instance_type)
            i += 1
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error searching for instances:")
    return instance_list
def cloudwatch_alarm():
    instance_list = []
    try:
        instance_list = instance_listing(['running'])   #function that takes in the instance status and returns an arrany of instance id's
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error searching for instances:")
    if not instance_list:  # check for an empty array i.e. no running instances
        print("No running instances")
        logging.info("No running instances")
        monitor_menu()
    else:
        choice = input("""Please select the instance number to monitor:""")
        print(choice)
        try:
            selectedinstance = instance_list[int(choice)]
            selectedinstance.monitor()
            alarm_name = "Webserver_NetworkIn_instance_" + selectedinstance.id
            cloudwatch_client = boto3.client('cloudwatch')
            #https://boto3.amazonaws.com/v1/documentation/api/latest/guide/cw-example-creating-alarms.html
            response = cloudwatch_client.put_metric_alarm(
                 AlarmName= alarm_name,
                 AlarmActions=['arn:aws:sns:eu-west-1:013355473762:test'],   #this arn will send an email to my email account
                 ComparisonOperator='LessThanThreshold',
                 EvaluationPeriods=1,
                 MetricName='NetworkIn',
                 Namespace='AWS/EC2',
                 Period=60,
                 Statistic='Average',
                 Threshold=30000,
                 ActionsEnabled=True,
                 AlarmDescription='Alarm when Network In < 30k',   #very low for example for assignment
                 Dimensions=[
                     {
                     'Name': 'InstanceId',
                     'Value': selectedinstance.id
                     },
                 ],
                 Unit='None'  #important to set this correctly as leaving it to the incorrect value will cause errors displaying in cloudwatch alarms page
            )
            print(response)
            response2 = cloudwatch_client.describe_alarm_history(
                 AlarmName=alarm_name,
                 HistoryItemType='Action',
                 StartDate=datetime(2015, 1, 1),
                 EndDate=datetime(2022, 1, 1),
                 MaxRecords=55,
            )
            print(response2)
            logging.info(response2)
            input("\nPress Enter to continue...")
            monitor_menu()
        except Exception as e:
            print(e)
            logging.warning(e)
            print("Incorrect choice, please try again")
            input("\nPress Enter to continue...")
            monitor_menu()
def pushmonitoring():
    instance_list = []
    try:
        instance_list = instance_listing(['running'])
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error searching for instances:")

    if not instance_list:
        logging.warning("No running instances")
        print("No running instances")
        input("\nPress Enter to continue...")
        monitor_menu()
    else:
        choice = input("""Which instance would you like to setup a server on?""")
        print(choice)
        try:
            selected_instance = instance_list[int(choice)]
            ssh_text = "ssh -o StrictHostkeyChecking=no -i kp20201.pem ec2-user@"
            selected_instance.monitor() #enabled detailed monitoring
            subprocess.run("scp -i kp20201.pem credentials ec2-user@" + selected_instance.public_ip_address + ":.",
                                   shell=True)  # scp to copy the monitoring file to the webserver
            subprocess.run("scp -i kp20201.pem config ec2-user@" + selected_instance.public_ip_address + ":.",
                                   shell=True)  # scp to copy the monitoring file to the webserver
            subprocess.run(
                ssh_text + selected_instance.public_ip_address + " \ 'mkdir ~/.aws; mv credentials ~/.aws/credentials; mv config ~/.aws/config'",
                shell=True) #move the credentials and aws config files to the required folders
            #create a file call mem.sh
            f = open("mem.sh", "w+")
            f.write("#!/bin/bash\n")
            f.write("USEDMEMORY=$(free -m | awk 'NR==2{printf \"%.2f\t\", $3*100/$2 }')\n")
            f.write("TCP_CONN=$(netstat -an | wc -l)\n")
            f.write("TCP_CONN_PORT_80=$(netstat -an | grep 80 | wc -l)\n")
            f.write("TCP_CONN_PORT_3000=$(netstat -an | grep 3000 | wc -l)\n")
            f.write("USERS=$(uptime |awk '{ print $6 }')\n")
            f.write("IO_WAIT=$(iostat | awk 'NR==4 {print $5}')\n")
            f.write("instance_id=" + selected_instance.id+"\n")
            f.write("aws cloudwatch put-metric-data --metric-name memory-usage --dimensions Instance=$instance_id  --namespace \"Custom\" --value $USEDMEMORY \n")
            f.write("aws cloudwatch put-metric-data --metric-name Tcp_connections --dimensions Instance=$instance_id  --namespace \"Custom\" --value $TCP_CONN\n")
            f.write("aws cloudwatch put-metric-data --metric-name TCP_connection_on_port_80 --dimensions Instance=$instance_id  --namespace \"Custom\" --value $TCP_CONN_PORT_80\n")
            f.write("aws cloudwatch put-metric-data --metric-name TCP_connection_on_port_3000 --dimensions Instance=$instance_id  --namespace \"Custom\" --value $TCP_CONN_PORT_3000\n")
            f.write("aws cloudwatch put-metric-data --metric-name IO_WAIT --dimensions Instance=$instance_id --namespace \"Custom\" --value $IO_WAIT\n")
            f.close()
            subprocess.run("scp -i kp20201.pem mem.sh ec2-user@" + selected_instance.public_ip_address + ":.",
                                   shell=True)  # scp to copy the monitoring file to the webserver
            subprocess.run("scp -i kp20201.pem cron.sh ec2-user@" + selected_instance.public_ip_address + ":.",
                                   shell=True)  # scp to copy the cronjob file to the webserver
            subprocess.run(
                ssh_text + selected_instance.public_ip_address + " \ 'chmod +x mem.sh; chmod +x cron.sh; ./mem.sh;sudo ./cron.sh'",
                shell=True) #change the permissions on the mem.sh and cron files and run both

            input("\nPress Enter to continue...")
            monitor_menu()
        except:
            logging.warning("Issue with choice entry to select an instance")
            print("Incorrect choice, as server may not be fully loaded,  please try again thanks")
            input("\nPress Enter to continue...")
            monitor_menu()
def custom_monitoring():
    instance_list = []
    try:
        instance_list = instance_listing(['running'])   #function that takes in the instance status and returns an arrany of instance id's
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error searching for instances:")
    if not instance_list:  # check for an empty array i.e. no running instances
        print("No running instances")
        logging.info("No running instances")
        monitor_menu()
    else:
        choice = input("""Please select the instance number to monitor:""")

        try:
            selectedinstance = instance_list[int(choice)]
            selectedinstance.monitor()
            print(selectedinstance.id)
            cloudwatch_client = boto3.client('cloudwatch')
            response = cloudwatch_client.get_metric_statistics(Namespace='Custom',
                                                        MetricName='IO_WAIT',Dimensions=[{'Value': selectedinstance.id, 'Name':'Instance'}], StartTime = datetime.utcnow() - timedelta(minutes=5),   # 5 minutes ago
                                             EndTime=datetime.utcnow(),                              # now
                                             Period=300,                                             # 5 min intervals
                                             Statistics=['Average'])
            print ("Average IO wait time:", response['Datapoints'][0]['Average'], "seconds")
            response = cloudwatch_client.get_metric_statistics(Namespace='Custom',
                                                        MetricName='Tcp_connections',Dimensions=[{'Value': selectedinstance.id, 'Name':'Instance'}], StartTime = datetime.utcnow() - timedelta(minutes=5),   # 5 minutes ago
                                             EndTime=datetime.utcnow(),                              # now
                                             Period=300,                                             # 5 min intervals
                                             Statistics=['Average'])
            print ("Tcp_connections:", response['Datapoints'][0]['Average'], "connections")
            response = cloudwatch_client.get_metric_statistics(Namespace='Custom',
                                                        MetricName='TCP_connection_on_port_80',Dimensions=[{'Value': selectedinstance.id, 'Name':'Instance'}], StartTime = datetime.utcnow() - timedelta(minutes=5),   # 5 minutes ago
                                             EndTime=datetime.utcnow(),                              # now
                                             Period=300,                                             # 5 min intervals
                                             Statistics=['Average'])
            print ("Tcp_connections on port 80:", response['Datapoints'][0]['Average'], "connections")
            response = cloudwatch_client.get_metric_statistics(Namespace='Custom',
                                                        MetricName='TCP_connection_on_port_3000',Dimensions=[{'Value': selectedinstance.id, 'Name':'Instance'}], StartTime = datetime.utcnow() - timedelta(minutes=5),   # 5 minutes ago
                                             EndTime=datetime.utcnow(),                              # now
                                             Period=300,                                             # 5 min intervals
                                             Statistics=['Average'])
            print ("Tcp_connections on port 3000:", response['Datapoints'][0]['Average'], "connections")
            response = cloudwatch_client.get_metric_statistics(Namespace='Custom',
                                                        MetricName='memory-usage',Dimensions=[{'Value': selectedinstance.id, 'Name':'Instance'}], StartTime = datetime.utcnow() - timedelta(minutes=5),   # 5 minutes ago
                                             EndTime=datetime.utcnow(),                              # now
                                             Period=300,                                             # 5 min intervals
                                             Statistics=['Average'])
            print ("Memory-usage:", response['Datapoints'][0]['Average'], "percent")
            input("\nPress Enter to continue...")

            monitor_menu()
        except Exception as e:
            print(e)
            logging.warning("Issue with choice entry as incorrect setup for custom data")
            print("Issue with choice entry as incorrect setup for custom data")
            choice = input("\nPress Enter to continue...or R to repeat")
            if choice == "R" or choice == "r":
                custom_monitoring()
            else:
               monitor_menu()
def open_logfile():
    try:
        subprocess.call('nano log/aws_assignment.log', shell=True)
        input("\nPress Enter to continue...")
        main()

    except Exception as e:
        print(e)
        input("\nPress Enter to continue...")
        main()
def quitInstance():
    instance_list = []
    try:
        instance_list = instance_listing(['running'])   #function that takes in the instance status and returns an arrany of instance id's
    except:
        logging.warning("Couldn't create a list of instances")
        print("Error searching for instances:")
    if not instance_list:  # check for an empty array i.e. no running instances
        print("No running instances")
        logging.info("No running instances")
        instance_menu()
    else:
        choice = input("""Please select the instance number to terminate:""")
        try:
            selectedinstance = instance_list[
                int(choice)]  # access the object returned for the selected instance
            print("The instance to be terminated is the following: " + selectedinstance.id)
            print("Please be patient")
            logging.info("The instance to be terminated is the following: " + selectedinstance.id)
            response = selectedinstance.terminate()
            logging.info(response)
            print(response)  # Printing the instance that has been terminated
            input("\nPress Enter to continue...")
            instance_menu()
        except Exception as error:
            print(error)
            logging.warning("Issue with choice entry to terminate an instance")
            print("Incorrect choice, please try again")
            input("\nPress Enter to continue...")
            instance_menu()

main()


def ten_numbers_min_max(n): # Замечательные числа 2
    min_number = n[0]
    max_number = n[0]
    for i in n:
        if min_number > i:
            min_number = i
        if max_number < i:
            max_number = i

    print(f'максимальных - {n.count(max_number)}, минимальных - {n.count(min_number)}')

n = (1,2,5,3,5,9,5,32,1,32,6,-4,5,-4,5,4,32,21,0)
ten_numbers_min_max(n)


import boto3
from io import BytesIO
import arviz as az
import os
import pandas as pd
import requests
import s3fs
import tempfile
import numpy as np
from covid.storage import (
    get_covidtracking_csv_key,
    get_processed_covidtracking_key,
    get_overall_output_key,
    get_inference_data_key,
    get_state_output_key,
)
from covid.data import (
    get_raw_covidtracking_data,
    process_covidtracking_data,
    summarize_inference_data,
)
from covid.models.generative import GenerativeModel

s3 = boto3.Session(
    aws_access_key_id=os.environ.get("S3_ACCESS_KEY", None),
    aws_secret_access_key=os.environ.get("S3_SECRET_KEY", None),
).resource(service_name="s3", endpoint_url=os.environ.get("S3_URL", None))
fs = s3fs.S3FileSystem(
    anon=False,
    key=os.environ.get("S3_ACCESS_KEY", None),
    secret=os.environ.get("S3_SECRET_KEY", None),
    client_kwargs=dict(endpoint_url=os.environ.get("S3_URL", None)),
)
S3_BUCKET = os.environ.get("RTLIVE_S3_BUCKET")

def task_get_covidtracking_data(run_date: pd.Timestamp):
    """ Cache COVIDTracking daily data to s3 """
    data = get_raw_covidtracking_data()
    key = get_covidtracking_csv_key(run_date)
    with fs.open(f"{S3_BUCKET}/{key}", "w") as file:
        data.to_csv(file)

def task_process_covidtracking_data(run_date: pd.Timestamp):
    """ Process the COVIDTracking daily data into a form that works with our
        model. """
    key = get_covidtracking_csv_key(run_date)
    with fs.open(f"{S3_BUCKET}/{key}") as file:
        data = pd.read_csv(file, parse_dates=["date"])
    data = process_covidtracking_data(data, run_date)
    key = get_processed_covidtracking_key(run_date)
    with fs.open(f"{S3_BUCKET}/{key}", "w") as file:
        data.to_csv(file)

def task_run_model(country: str, region: str, run_date: pd.Timestamp):
    """ Run the Generative model for a given region on a given date, store
        inference data into S3. """
    key = get_processed_covidtracking_key(run_date)
    with fs.open(f"{S3_BUCKET}/{key}") as file:
        df = pd.read_csv(file, index_col=["region", "date"], parse_dates=["date"])

    model_input = df.xs(region)
    gm = GenerativeModel(region, model_input)
    gm.sample()

    inference_data = gm.inference_data

    # Ensure no divergences
    assert (
        gm.n_divergences == 0
    ), f"Model {region} had {gm.n_divergences} divergences, failing."

    # Ensure convergence
    R_HAT_LIMIT = 1.1
    r_hat = az.rhat(inference_data).to_dataframe().fillna(1.0)
    assert r_hat.le(R_HAT_LIMIT).all().all(), f"r_hat exceeded threshold, failing."

    with tempfile.NamedTemporaryFile() as fp:
        inference_data.to_netcdf(fp.name)
        fp.seek(0)
        s3.Bucket(S3_BUCKET).upload_fileobj(
            fp, get_inference_data_key(run_date, region, country=country)
        )
    return {"country": country, "region": region, "r_hat": r_hat}

def task_render_region_result(country: str, region: str, run_date: pd.Timestamp):
    """ Render a CSV with summary output for a given region """
    az.rcParams["data.load"] = "eager"

    with tempfile.NamedTemporaryFile() as fp:
        s3.Bucket(S3_BUCKET).download_file(
            get_inference_data_key(run_date, region, country=country), fp.name
        )
        fp.seek(0)
        inference_data = az.from_netcdf(fp.name)

    summary = summarize_inference_data(inference_data)
    key = get_state_output_key(run_date, region, country=country)
    with fs.open(f"{S3_BUCKET}/{key}", "w") as file:
        summary.to_csv(file)

def task_gather_region_results(country: str, regions: [str], run_date: pd.Timestamp):
    """ Collects all regions results and outputs them as a single file """
    dfs = []
    for region in regions:
        key = get_state_output_key(run_date, region, country=country)
        with fs.open(f"{S3_BUCKET}/{key}") as file:
            df = pd.read_csv(file)
        multi_index = pd.MultiIndex.from_product([[region], df.index])
        multi_index.set_names(["region", "index"], inplace=True)
        df.index = multi_index
        dfs.append(df)
    all_regions = pd.concat(dfs)
    all_key = get_overall_output_key(run_date, country=country)
    with fs.open(f"{S3_BUCKET}/{all_key}", "w") as file:
        all_regions.to_csv(file)


import sys
import utils

def findNearestPoint(points, used, src):
    # If no nearest point, return max.
    dest = src
    minDist = sys.float_info.max
    i = (src + 1) % len(points)
    for j in range(len(points) - 1):
        if not used[i]:
            dist = utils.length(points[src], points[i])
            if dist < minDist:
                dest = i
                minDist = dist
        i = (i + 1) % len(points)
    return dest, minDist

def solve(points):
    tour = [0 for i in range(len(points))]
    used = [False for i in range(len(points))]
    totalDist = 0.0
    used[0] = True
    src = 0
    for i in range(1, len(points)):
        dest, minDist = findNearestPoint(points, used, src)
        tour[i] = dest
        used[dest] = True
        src = dest
        totalDist += minDist
    return 0, totalDist + utils.length(points[tour[-1]], points[tour[0]]), tour


from nltk.tree import Tree

from allennlp.common.testing import AllenNlpTestCase
from allennlp.data.dataset_readers.dataset_utils.span_utils import enumerate_spans

from allennlp_models.structured_prediction.dataset_readers.penn_tree_bank import (
    PennTreeBankConstituencySpanDatasetReader,
)
from tests import FIXTURES_ROOT

class TestPennTreeBankConstituencySpanReader(AllenNlpTestCase):
    def setup_method(self):
        super().setup_method()
        self.span_width = 5

    def test_read_from_file(self):

        ptb_reader = PennTreeBankConstituencySpanDatasetReader()
        instances = list(
            ptb_reader.read(str(FIXTURES_ROOT / "structured_prediction" / "example_ptb.trees"))
        )

        assert len(instances) == 2

        fields = instances[0].fields
        tokens = [x.text for x in fields["tokens"].tokens]
        pos_tags = fields["pos_tags"].labels
        spans = [(x.span_start, x.span_end) for x in fields["spans"].field_list]
        span_labels = fields["span_labels"].labels

        assert tokens == [
            "Also",
            ",",
            "because",
            "UAL",
            "Chairman",
            "Stephen",
            "Wolf",
            "and",
            "other",
            "UAL",
            "executives",
            "have",
            "joined",
            "the",
            "pilots",
            "'",
            "bid",
            ",",
            "the",
            "board",
            "might",
            "be",
            "forced",
            "to",
            "exclude",
            "him",
            "from",
            "its",
            "deliberations",
            "in",
            "order",
            "to",
            "be",
            "fair",
            "to",
            "other",
            "bidders",
            ".",
        ]
        assert pos_tags == [
            "RB",
            ",",
            "IN",
            "NNP",
            "NNP",
            "NNP",
            "NNP",
            "CC",
            "JJ",
            "NNP",
            "NNS",
            "VBP",
            "VBN",
            "DT",
            "NNS",
            "POS",
            "NN",
            ",",
            "DT",
            "NN",
            "MD",
            "VB",
            "VBN",
            "TO",
            "VB",
            "PRP",
            "IN",
            "PRP$",
            "NNS",
            "IN",
            "NN",
            "TO",
            "VB",
            "JJ",
            "TO",
            "JJ",
            "NNS",
            ".",
        ]

        assert spans == enumerate_spans(tokens)
        gold_tree = Tree.fromstring(
            "(S(ADVP(RB Also))(, ,)(SBAR(IN because)"
            "(S(NP(NP(NNP UAL)(NNP Chairman)(NNP Stephen)(NNP Wolf))"
            "(CC and)(NP(JJ other)(NNP UAL)(NNS executives)))(VP(VBP have)"
            "(VP(VBN joined)(NP(NP(DT the)(NNS pilots)(POS '))(NN bid))))))"
            "(, ,)(NP(DT the)(NN board))(VP(MD might)(VP(VB be)(VP(VBN "
            "forced)(S(VP(TO to)(VP(VB exclude)(NP(PRP him))(PP(IN from)"
            "(NP(PRP$ its)(NNS deliberations)))(SBAR(IN in)(NN order)(S("
            "VP(TO to)(VP(VB be)(ADJP(JJ fair)(PP(TO to)(NP(JJ other)(NNS "
            "bidders))))))))))))))(. .))"
        )

        assert fields["metadata"].metadata["gold_tree"] == gold_tree
        assert fields["metadata"].metadata["tokens"] == tokens

        correct_spans_and_labels = {}
        ptb_reader._get_gold_spans(gold_tree, 0, correct_spans_and_labels)
        for span, label in zip(spans, span_labels):
            if label != "NO-LABEL":
                assert correct_spans_and_labels[span] == label

        fields = instances[1].fields
        tokens = [x.text for x in fields["tokens"].tokens]
        pos_tags = fields["pos_tags"].labels
        spans = [(x.span_start, x.span_end) for x in fields["spans"].field_list]
        span_labels = fields["span_labels"].labels

        assert tokens == [
            "That",
            "could",
            "cost",
            "him",
            "the",
            "chance",
            "to",
            "influence",
            "the",
            "outcome",
            "and",
            "perhaps",
            "join",
            "the",
            "winning",
            "bidder",
            ".",
        ]

        assert pos_tags == [
            "DT",
            "MD",
            "VB",
            "PRP",
            "DT",
            "NN",
            "TO",
            "VB",
            "DT",
            "NN",
            "CC",
            "RB",
            "VB",
            "DT",
            "VBG",
            "NN",
            ".",
        ]

        assert spans == enumerate_spans(tokens)

        gold_tree = Tree.fromstring(
            "(S(NP(DT That))(VP(MD could)(VP(VB cost)(NP(PRP him))"
            "(NP(DT the)(NN chance)(S(VP(TO to)(VP(VP(VB influence)(NP(DT the)"
            "(NN outcome)))(CC and)(VP(ADVP(RB perhaps))(VB join)(NP(DT the)"
            "(VBG winning)(NN bidder)))))))))(. .))"
        )

        assert fields["metadata"].metadata["gold_tree"] == gold_tree
        assert fields["metadata"].metadata["tokens"] == tokens

        correct_spans_and_labels = {}
        ptb_reader._get_gold_spans(gold_tree, 0, correct_spans_and_labels)
        for span, label in zip(spans, span_labels):
            if label != "NO-LABEL":
                assert correct_spans_and_labels[span] == label

    def test_strip_functional_tags(self):
        ptb_reader = PennTreeBankConstituencySpanDatasetReader()
        # Get gold spans should strip off all the functional tags.
        tree = Tree.fromstring(
            "(S (NP=PRP (D the) (N dog)) (VP-0 (V chased) (NP|FUN-TAGS (D the) (N cat))))"
        )
        ptb_reader._strip_functional_tags(tree)
        assert tree == Tree.fromstring(
            "(S (NP (D the) (N dog)) (VP (V chased) (NP (D the) (N cat))))"
        )

    def test_get_gold_spans_correctly_extracts_spans(self):
        ptb_reader = PennTreeBankConstituencySpanDatasetReader()
        tree = Tree.fromstring("(S (NP (D the) (N dog)) (VP (V chased) (NP (D the) (N cat))))")

        span_dict = {}
        ptb_reader._get_gold_spans(tree, 0, span_dict)
        spans = list(span_dict.items())
        assert spans == [((0, 1), "NP"), ((3, 4), "NP"), ((2, 4), "VP"), ((0, 4), "S")]

    def test_get_gold_spans_correctly_extracts_spans_with_nested_labels(self):
        ptb_reader = PennTreeBankConstituencySpanDatasetReader()
        # Here we have a parse with several nested labels - particularly the (WHNP (WHNP (WP What)))
        # fragment. These should be concatenated into a single label by get_gold_spans.
        tree = Tree.fromstring(
            """
            (S
        (` ``)
        (S-TPC
        (NP-SBJ (PRP We))
        (VP
            (VBP have)
            (S
            (VP
                (TO to)
                (VP
                (VP
                    (VB clear)
                    (PRT (RP up))
                    (NP (DT these) (NNS issues)))
                (CC and)
                (VP
                    (VB find)
                    (PRT (RP out))
                    (SBAR-NOM
                    (WHNP (WHNP (WP what)))
                    (S
                        (VP
                        (VBZ is)
                        (ADJP-PRD (JJ present))
                        (SBAR
                            (WHNP (WDT that))
                            (S
                            (VP
                                (VBZ is)
                                (VP
                                (VBG creating)
                                (NP (JJ artificial) (NN volatility)))))))))))))))
        (, ,)
        ('' '')
        (NP-SBJ (NNP Mr.) (NNP Fisher))
        (VP (VBD said))
        (. .))
        """
        )
        span_dict = {}
        ptb_reader._strip_functional_tags(tree)
        ptb_reader._get_gold_spans(tree, 0, span_dict)
        assert span_dict == {
            (1, 1): "NP",
            (5, 5): "PRT",
            (6, 7): "NP",
            (4, 7): "VP",
            (10, 10): "PRT",
            (11, 11): "WHNP-WHNP",
            (13, 13): "ADJP",
            (14, 14): "WHNP",
            (17, 18): "NP",
            (16, 18): "VP",
            (15, 18): "S-VP",
            (14, 18): "SBAR",
            (12, 18): "S-VP",
            (11, 18): "SBAR",
            (9, 18): "VP",
            (4, 18): "VP",
            (3, 18): "S-VP",
            (2, 18): "VP",
            (1, 18): "S",
            (21, 22): "NP",
            (23, 23): "VP",
            (0, 24): "S",
        }


"""
Phil Boynton
bskine@gmail.com

T1 diabetic blood glucose graphing script for use with Dexcom6
"""

import averages

from bokeh.io import output_file
from bokeh.plotting import figure, show
from bokeh.transform import factor_cmap
from bokeh.models import ColumnDataSource
from bokeh.palettes import Spectral6

# from bokeh.layouts import row, column, gridplot
# import bokeh.models.widgets

output_file("average_bar_graph.html", title="Average Blood Glucose")

top = [
    averages.glucose_current()[1],
    averages.glucose_average_1hour()[1],
    averages.glucose_average_3hour()[1],
    averages.glucose_average_6hour()[1],
    averages.glucose_average_12hour()[1],
    averages.glucose_average_24hour()[1],
]

tags = ["current", "1Hour", "3hour", "6Hour", "12Hour", "24Hour"]
x = [tags[i] for i in range(0, 6)]

source = ColumnDataSource(data=dict(x=x, top=top))

fig = figure(
    title="CURRENT AND AVERAGE BLOOD GLUCOSE",
    height=500,
    width=500,
    x_range=x,
    y_range=(0, 300),
    x_axis_location="below",
    x_axis_label="hourly average",
    y_axis_label="mg/dl",
    toolbar_location=None,
    tools=""
)

fig.vbar(
    x,
    top=top,
    # color='blue',
    fill_color=factor_cmap("x", palette=Spectral6, factors=x),
    width=0.75,
)

fig.line(x=x, y=110, color='black', dash='dashed')

show(fig)


import ctypes

_at_err_dict = {
    1: 'NOTINITIALISED',
    2: 'NOTIMPLEMENTED',
    3: 'READONLY',
    4: 'NOTREADABLE',
    5: 'NOTWRITABLE',
    6: 'OUTOFRANGE',
    7: 'INDEXNOTAVAILABLE',
    8: 'INDEXNOTIMPLEMENTED',
    9: 'EXCEEDEDMAXSTRINGLENGTH',
    10: 'CONNECTION',
    11: 'NODATA',
    12: 'INVALIDHANDLE',
    13: 'TIMEDOUT',
    14: 'BUFFERFULL',
    15: 'INVALIDSIZE',
    16: 'INVALIDALIGNMENT',
    17: 'COMM',
    18: 'STRINGNOTAVAILABLE',
    19: 'STRINGNOTIMPLEMENTED',
    20: 'NULL_FEATURE',
    21: 'NULL_HANDLE',
    22: 'NULL_IMPLEMENTED_VAR',
    23: 'NULL_READABLE_VAR',
    24: 'NULL_READONLY_VAR',
    25: 'NULL_WRITABLE_VAR',
    26: 'NULL_MINVALUE',
    27: 'NULL_MAXVALUE',
    28: 'NULL_VALUE',
    29: 'NULL_STRING',
    30: 'NULL_COUNT_VAR',
    31: 'NULL_ISAVAILABLE_VAR',
    32: 'NULL_MAXSTRINGLENGTH',
    33: 'NULL_EVCALLBACK',
    34: 'NULL_QUEUE_PTR',
    35: 'NULL_WAIT_PTR',
    36: 'NULL_PTRSIZE',
    37: 'NOMEMORY',
    38: 'DEVICEINUSE',
    100: 'HARDWARE_OVERFLOW',
    1002: 'AT_ERR_INVALIDOUTPUTPIXELENCODING',
    1003: 'AT_ERR_INVALIDINPUTPIXELENCODING',
    1004: 'AT_ERR_INVALIDMETADATAINFO'
}

class AndorError(RuntimeError):
    pass

def _at_errcheck(result, func, args):
    if result != 0:
        err_type = _at_err_dict[result]
        func_name = func.__name__
        arg_str = ' ,'.join(map(str, args))
#       with open('/home/ehvatum/d.txt', 'a') as f:
#           print('{} error when calling {}({})'.format(err_type, func_name, args), file=f)
        raise AndorError('{} error when calling {}({})'.format(err_type, func_name, args))
    return args

# NB: Callbacks should return AT_CALLBACK_SUCCESS, or, equivalently, 0
FeatureCallback = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_wchar_p, ctypes.c_void_p)

ANDOR_INFINITE = 0xFFFFFFFF
AT_FALSE = 0
AT_CALLBACK_SUCCESS = 0


import numpy as np
import h5py

if __name__ == '__main__':
    f = h5py.File('./Groningen_Data_Example.hdf5', 'r')
    # keys() shows the contents of the file

    # extract all the data.
    # note that these hdf4 datasets should behave exactly as np arrays

    # The reservoir data
    x_res = f['x_res']
    y_res = f['y_res']
    reservoir_outline = f['res_outline']
    x_outline = f['x_outline']
    y_outline = f['y_outline']
    # The original well extraction data
    extractions = f['extraction_data']
    wells_locations = f['wells_locations_list']
    wells_names = f['wells_names']
    # The computed pressures from the linearized pressure diffusion model
    pressures = f['pressures_list']
    x_mesh = f['x_mesh']  # The meshes used for the pressure diffusion model
    y_mesh = f['y_mesh']  # The meshes used for the pressure diffusion model
    # The computed deformations and coulomb stress change from the BorisBlocks model
    deformations = f['deformations_list']
    # The meshes used for the BorisBlocks model
    x_reservoir_mesh = f['x_reservoir_mesh']
    # The meshes used for the BorisBlocks model
    y_reservoir_mesh = f['y_reservoir_mesh']
    max_coulomb_stress = f['max_coulomb_stress']
    # the maximum coulomb stress smoothed with a 6km Gaussian2D kernel calculated 10m below the reservoir.
    smoothed_coulomb_stress = f['smoothed_coulomb_stress']

    # The simulation given in this example is from 1956 to 2019 for a total of 756 months
    y0, y1, dy = 1956, 2019, 12

    dates = np.linspace(y0, y1+1, (y1-y0)*dy+1)


import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from sklearn.preprocessing import LabelEncoder

from typing import List, Dict

"""
Constants
"""

COMPRESSION_GZIP = "gzip"
DIAGNOSIS_DATA_PATH = "data/MIMIC3_DIAGNOSES_ICD_subset.csv.gz"
DESC_DATA_PATH = "data/D_ICD_DIAGNOSES.csv.gz"

COLUMN_SUBJECT_ID = "SUBJECT_ID"
COLUMN_ICD9_CODE = "ICD9_CODE"
COLUMN_SHORT_TITLE = "SHORT_TITLE"

K = 5  # Number of topics
TOTAL_EPOCHS = 100
TOP_ICD_COUNT = 10
TOP_PATIENT_COUNT = 100

ALPHA = 1  # Document topic Dirichlet prior
BETA = 0.001  # ICD-9 topic Dirichlet prior

"""
ICD Code Data Processor
"""

class ICDCodeDataProcessor:

    """ Public methods """

    def create_patient_docs(self, df: pd.DataFrame) -> Dict[int, List[int]]:
        """
        @param self:
        @param df: encoded diagnosis dataframe
        @return: patient docs
        """

        result: Dict[int, List[int]] = None

        if df is not None:

            result = dict()
            for _, row in df.iterrows():

                patient_id = row[COLUMN_SUBJECT_ID]
                doc = result.get(patient_id)
                if doc is None:
                    doc: List[int] = list()

                icd_code = row[COLUMN_ICD9_CODE]
                doc.append(icd_code)
                result[patient_id] = doc

        return result

"""
Plot
"""

class ICDPlot:

    """ Public methods """

    def plot_top_icds(self, descriptions: pd.DataFrame,
                      phi: np.array, topic_labels: List[str],
                      icd_encoder: LabelEncoder) -> None:

        top_icds = np.zeros((K * TOP_ICD_COUNT, K))
        top_icds_labels = []
        for k in range(K):

            # Top ICD codes
            top_idx = np.argsort(phi[:, k])[::-1][:TOP_ICD_COUNT]
            top_value = np.take(phi[:, k], top_idx)
            top_icds[k * TOP_ICD_COUNT: (k + 1) * TOP_ICD_COUNT, k] = top_value

            # Top labels
            icd_codes = icd_encoder.inverse_transform(top_idx)
            if icd_codes is not None and icd_codes.shape[0] > 0:

                for icd_code in icd_codes:

                    icd_code_value = str(icd_code)
                    description_row = descriptions.loc[descriptions[COLUMN_ICD9_CODE].str.startswith(
                        icd_code_value)]

                    if description_row.empty:
                        continue

                    description_row = description_row.iloc[0]
                    short_title = description_row[COLUMN_SHORT_TITLE]
                    label = icd_code_value + "-" + short_title

                    top_icds_labels.append(label)

        assert len(top_icds_labels) == K * TOP_ICD_COUNT

        plt.clf()
        plt.figure(figsize=(3, 10))
        plt.title("Latent topics inferred from the collapsed Gibbs sampling")
        plt.tight_layout()
        axs = sns.heatmap(top_icds, cmap="Reds",
                          yticklabels=top_icds_labels,
                          xticklabels=topic_labels)
        cbar = axs.collections[0].colorbar
        cbar.ax.set_title("topic prob", fontsize=9)
        plt.savefig("figure_1_top_icds.png", bbox_inches='tight')
        plt.close()

    def plot_correlation(self, D: int, theta: np.array,
                         patient_docs: Dict[int, List[int]],
                         topic_labels: List[str],
                         icd_xticklabels: List[str],
                         icd_encoder: LabelEncoder) -> np.array:

        target_icd_codes = np.zeros((D, 3))
        for d, doc in enumerate(patient_docs.values()):

            icd_codes = icd_encoder.inverse_transform(doc)
            if icd_codes is not None and icd_codes.shape[0] > 0:

                for icd_code in icd_codes:

                    icd_code_value = str(icd_code)
                    if icd_code_value.startswith("331"):
                        target_icd_codes[d, 0] = 1

                    elif icd_code_value.startswith("332"):
                        target_icd_codes[d, 1] = 1

                    elif icd_code_value.startswith("340"):
                        target_icd_codes[d, 2] = 1

        total_correlation = np.corrcoef(theta.T, target_icd_codes.T)  # 8 x 8

        # First 5 rows are correlation from 5 topics
        # Last 3 columns are correlation from 3 ICD codes
        correlation = total_correlation[:5, -3:]

        plt.clf()
        plt.figure(figsize=(6, 5))
        plt.title("Topics by target ICD correlation")
        plt.tight_layout()
        axs = sns.heatmap(correlation, cmap="coolwarm",
                          yticklabels=topic_labels,
                          xticklabels=icd_xticklabels)
        cbar = axs.collections[0].colorbar
        cbar.ax.set_title("ICD-topic cor", fontsize=9)
        plt.yticks(rotation=0)
        plt.savefig("figure_2_correlation.png", bbox_inches='tight')
        plt.close()

        return target_icd_codes

    def plot_top_patients(self, theta: np.array, target_icd_codes: np.array,
                          topic_labels: List[str],
                          icd_xticklabels: List[str]) -> None:

        top_patients_idx = np.zeros((K * TOP_PATIENT_COUNT)).astype(int)
        for k in range(K):
            # Top patients idx
            top_idx = np.argsort(theta[:, k])[::-1][:TOP_PATIENT_COUNT]
            top_patients_idx[k *
                             TOP_PATIENT_COUNT: (k + 1) * TOP_PATIENT_COUNT] = top_idx

        top_patients = theta[top_patients_idx]
        top_target_icd_codes = target_icd_codes[top_patients_idx]

        cmap_1 = LinearSegmentedColormap.from_list(
            "custom", ["#FFFFFF", "#000000"])
        cmap_2 = LinearSegmentedColormap.from_list(
            "custom", ["#FFFFFF", "#FF0000"])
        cmap_3 = LinearSegmentedColormap.from_list(
            "custom", ["#FFFFFF", "#0000FF"])
        cmaps = [cmap_1, cmap_2, cmap_3]
        plt.clf()
        fig, axs = plt.subplots(nrows=1, ncols=top_target_icd_codes.shape[1] + 1,
                                width_ratios=[1, 1, 1, 12],
                                gridspec_kw={'wspace': 0},
                                figsize=(6, 10))
        fig.suptitle("Top 100 patients per topic")
        annotation_axs = axs[:3]
        for i, (a, c, l) in enumerate(zip(annotation_axs, cmaps, icd_xticklabels)):
            sns.heatmap(top_target_icd_codes[:, i][:, None],
                        ax=a, cmap=c,
                        xticklabels=[l],
                        yticklabels=False,
                        cbar=False)
            a.tick_params(labelrotation=90)

        top_patients_axs = axs[-1]
        sns.heatmap(top_patients, cmap="Reds",
                    ax=top_patients_axs,
                    yticklabels=False,
                    xticklabels=topic_labels)
        top_patients_axs.tick_params(labelrotation=90)
        cbar = top_patients_axs.collections[0].colorbar
        cbar.ax.set_title("topic prob", fontsize=9)
        fig.savefig("figure_3_top_patients.png", bbox_inches='tight')
        plt.close()

"""
LDA
"""

class LDAModel:

    """ Initialize """

    def __init__(self, K: int, D: int, M: int,
                 alpha: float, beta: float,
                 patient_docs: Dict[int, List[int]]) -> None:
        """
        @param self:
        @param K: Number of topics
        @param D: Number of unique patients (documents)
        @param M: Number of unique ICD codes (vocabulary)
        @param alpha: Hyper-parameter alpha to generate topic distribution
        @param beta: Hyper-parameter beta to generate word x topic matrix
        @param patient_docs: patient docs
        """

        self._K = K
        self._D = D
        self._M = M
        self._alpha = alpha
        self._beta = beta
        self._patient_docs = patient_docs

        self._initialize()

    """ Getters """

    @property
    def theta(self) -> np.array:
        if self._theta is None:
            raise Exception("Please run fit() first!")
        return self._theta

    @property
    def phi(self) -> np.array:
        if self._phi is None:
            raise Exception("Please run fit() first!")
        return self._phi

    """ Public methods """

    def fit(self, epochs: int) -> None:
        """
        @param self:
        @param epochs: Training epochs
        """

        for _ in range(epochs):

            for d, doc in enumerate(self._patient_docs.values()):
                for m in doc:

                    # Sample new topic
                    updated_topic = self._sample_topic(d, m)

                    # Update count
                    self._z[d, m] = updated_topic
                    self._doc_topic[d, updated_topic] += 1
                    self._topic_terms[updated_topic, m] += 1
                    self._terms_count_in_topic[updated_topic] += 1

        # Update theta and phi
        self._update_parameter()

    """ Private methods """

    def _initialize(self) -> None:

        # Topics
        self._z = np.zeros((self._D, self._M)).astype(
            int)  # Document x Vocabulary

        # The number of words assigned to topic k in document d
        self._doc_topic = np.zeros((self._D, self._K))  # Document x Topic

        # The number of words in document d
        self._terms_count_in_doc = np.zeros((self._D))  # Document x 1

        # The number of word m assigned to topic k in all documents
        self._topic_terms = np.zeros((self._K, self._M))  # Topic x Vocabulary

        # The number of words assigned to topic k in all documents
        self._terms_count_in_topic = np.zeros((self._K))  # Topics x 1

        # Topic distribution, Document x Topic
        self._theta = np.zeros((self._D, self._K))

        # Word x Topic
        self._phi = np.zeros((self._M, self._K))

        # Initialize
        for d, doc in enumerate(self._patient_docs.values()):

            for m in doc:

                rand_topic = np.random.randint(self._K)

                self._z[d, m] = rand_topic
                self._doc_topic[d, rand_topic] += 1
                self._topic_terms[rand_topic, m] += 1
                self._terms_count_in_topic[rand_topic] += 1

            # Total number of words in document d
            self._terms_count_in_doc[d] = len(doc)

    def _sample_topic(self, d: int, m: int) -> int:

        # No need to consider the current word
        old_topic = self._z[d, m]
        self._doc_topic[d, old_topic] -= 1
        self._topic_terms[old_topic, m] -= 1
        self._terms_count_in_topic[old_topic] -= 1

        # Compute p(z=k)
        probs = np.zeros((self._K))
        for k in range(self._K):

            first_term = self._alpha + self._doc_topic[d, k]
            numerator = self._beta + self._topic_terms[k, m]
            denominator = self._M * self._beta + self._terms_count_in_topic[k]
            second_terms = numerator / denominator
            probs[k] = first_term * second_terms

        # Normalize p(z=k)
        probs = probs / np.sum(probs)

        # Sample new topic
        sample = np.random.multinomial(1, probs)
        return np.nonzero(sample)[0][0]

    def _update_parameter(self) -> None:

        # Update topic distribution (theta)
        for d in range(len(self._patient_docs)):
            for k in range(self._K):

                numerator = self._alpha + self._doc_topic[d, k]
                denominator = self._K * self._alpha + \
                    self._terms_count_in_doc[d]

                self._theta[d, k] = numerator / denominator

        # Update word x topic matrix (phi)
        for m in range(self._M):
            for k in range(self._K):

                numerator = self._beta + self._topic_terms[k, m]
                denominator = self._M * self._beta + \
                    self._terms_count_in_topic[k]

                self._phi[m, k] = numerator / denominator

"""
Main
"""

def main():

    # Load data
    diagnosis = pd.read_csv(DIAGNOSIS_DATA_PATH,
                            compression=COMPRESSION_GZIP)
    print("Diagnosis shape: {} \n{}".format(diagnosis.shape, diagnosis.head()))

    descriptions = pd.read_csv(DESC_DATA_PATH,
                               compression=COMPRESSION_GZIP)
    print("Description shape: {} \n{}".format(
        descriptions.shape, descriptions.head()))

    D = len(diagnosis[COLUMN_SUBJECT_ID].unique())
    M = len(diagnosis[COLUMN_ICD9_CODE].unique())

    # Preprocess data
    patient_encoder = LabelEncoder()
    icd_encoder = LabelEncoder()

    encoded_patients = patient_encoder.fit_transform(
        diagnosis[COLUMN_SUBJECT_ID])
    encoded_icds = icd_encoder.fit_transform(diagnosis[COLUMN_ICD9_CODE])

    encoded_diagnosis = pd.DataFrame(
        columns=[COLUMN_SUBJECT_ID, COLUMN_ICD9_CODE])
    encoded_diagnosis[COLUMN_SUBJECT_ID] = encoded_patients
    encoded_diagnosis[COLUMN_ICD9_CODE] = encoded_icds

    data_processor = ICDCodeDataProcessor()
    patient_docs = data_processor.create_patient_docs(encoded_diagnosis)

    # Q1: Fit LDA model
    lda = LDAModel(K=K, D=D, M=M,
                   alpha=ALPHA, beta=BETA,
                   patient_docs=patient_docs)
    lda.fit(epochs=TOTAL_EPOCHS)

    phi = lda.phi
    theta = lda.theta

    # Q2: Visualise the top ICD codes under each topic
    icd_plot = ICDPlot()
    topic_labels = ["topic {}".format(k + 1) for k in range(K)]
    icd_plot.plot_top_icds(descriptions, phi, topic_labels, icd_encoder)

    # Q3: Correlating topics with the target ICD codes
    icd_xticklabels = ["icd331", "icd332", "icd340"]
    target_icd_codes = icd_plot.plot_correlation(D, theta, patient_docs,
                                                 topic_labels, icd_xticklabels, icd_encoder)

    # Q4: Visualizing patient topic mixtures
    icd_plot.plot_top_patients(theta, target_icd_codes,
                               topic_labels, icd_xticklabels)

if __name__ == "__main__":
    main()


import os
import copy
import glob
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from transformers import Trainer
from transformers.trainer_pt_utils import nested_detach
from openprompt.utils.reproduciblity import set_seed
from openprompt import PromptDataLoader, PromptForClassification
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
from datasets import load_dataset, load_metric
from openprompt.plms import load_plm
from openprompt.prompts import SoftTemplate, ManualVerbalizer

from .data_processor import data_processor_list
from .utils import decorate

class PromptHub(Trainer):
    def __init__(self, prompt_emb=None, **kwargs):
        args = kwargs['args']
        self.out_dir_root = args.output_dir

        processor = data_processor_list[args.dataset]()

        # Model
        template_text = '{"soft": None, "duplicate": ' + str(args.prompt_len) + ', "same": True} {"mask"} {"placeholder": "text_a"} {"placeholder": "text_b"}'

        model, template, verbalizer, plm, tokenizer, model_config, tokenizer_wrapper_class, model_type = self.get_model(args.backbone, processor, args)

        self.set_active_state_dict(model)   # Only save soft prompts

        # Initialize transformers.trainer
        kwargs['model'] = model
        kwargs['tokenizer'] = tokenizer
        kwargs['train_dataset'] = processor.train_dataset
        kwargs['eval_dataset'] = processor.eval_dataset
        super().__init__(**kwargs)

        self.config = model_config
        self.plm = plm
        self.template_text = template_text
        self.template = template
        self.verbalizer = verbalizer
        self.tokenizer_wrapper_class = tokenizer_wrapper_class
        self.prompt_emb = None

        # Transfer
        self.source_model_type = model_type
        self.target_model_type = None

        print('Trainable parameters:')
        for n, p in self.model.named_parameters():
            if p.requires_grad:
                print(n, p.shape)

        print(f'Template: {self.template}')
        print(f'Verbalizer: {self.verbalizer}')
        print(f'Raw input example: {self.train_dataset[0]}')
        print(f'Wrapped input example: {self.template.wrap_one_example(self.train_dataset[0])}')

    def get_model(self, model_name, processor, args):
        if 'bert-' in model_name:
            model_type = 'bert'

        elif 'roberta-' in model_name:
            model_type = 'roberta'

        elif 't5-' in model_name:
            model_type = 't5'

        elif 'gpt' in model_name:
            model_type = 'gpt'

        # Load openprompt models
        if (not hasattr(self, 'args')) or args.backbone != model_name:
            plm, tokenizer, model_config, tokenizer_wrapper_class = load_plm(model_type, model_name)
        else:
            plm = self.plm
            tokenizer = self.tokenizer
            model_config = self.config
            tokenizer_wrapper_class = self.tokenizer_wrapper_class
            model_type = self.source_model_type

        # Load soft template
        if hasattr(self, 'template') and self.template is not None:
            template = self.template
        else:
            template_text = '{"soft": None, "duplicate": ' + str(args.prompt_len) + ', "same": True} {"mask"} {"placeholder": "text_a"} {"placeholder": "text_b"}'
            template = SoftTemplate(model=plm, text=template_text, tokenizer=tokenizer, num_tokens=args.prompt_len) # initialize_from_vocab=args.init_from_vocab

        # Load verbalizer
        if hasattr(self, 'verbalizer') and self.verbalizer is not None:
            verbalizer = self.verbalizer
        else:
            verbalizer = ManualVerbalizer(tokenizer, classes=processor.labels, label_words=processor.label_words)

        if hasattr(self, 'model') and self.model is not None:
            model = self.model
        else:
            model = PromptForClassification(plm=plm, template=template, verbalizer=verbalizer, freeze_plm=True)

            if hasattr(args, "model_parallel") and args.model_parallel:
                print('parallelize model!')
                model.parallelize()

        _keys_to_ignore_on_save = []
        for n, p in model.named_parameters():
            if not p.requires_grad:
                _keys_to_ignore_on_save.append(n)

        model._keys_to_ignore_on_save = _keys_to_ignore_on_save

        return model, template, verbalizer, plm, tokenizer, model_config, tokenizer_wrapper_class, model_type

    def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(inputs)
        # Save past state if it exists
        # TODO: this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        loss = nn.CrossEntropyLoss()(outputs, inputs['label'])

        return (loss, outputs) if return_outputs else loss

    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys):
        inputs = self._prepare_inputs(inputs)
        if ignore_keys is None:
            if hasattr(self.model, "config"):
                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
            else:
                ignore_keys = []

        # labels may be popped when computing the loss (label smoothing for instance) so we grab them first.
        labels = nested_detach(tuple(inputs.__getitem__(name) for name in ['label']))
        if len(labels) == 1:
            labels = labels[0]

        with torch.no_grad():
            with self.compute_loss_context_manager():
                loss, logits = self.compute_loss(model, inputs, return_outputs=True)
            loss = loss.mean().detach()

        if prediction_loss_only:
            return (loss, None, None)

        logits = nested_detach(logits)
        if len(logits) == 1:
            logits = logits[0]

        return (loss, logits, labels)

    def _prepare_input(self, data):
        kwargs = dict(device=self.args.device)
        return data.to(**kwargs)

    def _prepare_inputs(self, inputs):
        inputs = self._prepare_input(inputs)
        return inputs

    def get_train_dataloader(self) -> DataLoader:
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        train_dataloader = PromptDataLoader(
            dataset=self.train_dataset,
            template=self.template,
            tokenizer=self.tokenizer,
            tokenizer_wrapper_class=self.tokenizer_wrapper_class,
            batch_size=self._train_batch_size,
            max_seq_length=self.args.max_source_length,
            decoder_max_length=1,
            shuffle=True)

        return train_dataloader

    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")

        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset

        validation_dataloader = PromptDataLoader(
            dataset=eval_dataset,
            template=self.template,
            tokenizer=self.tokenizer,
            tokenizer_wrapper_class=self.tokenizer_wrapper_class,
            batch_size=self.args.per_device_eval_batch_size,
            max_seq_length=self.args.max_source_length,
            decoder_max_length=1,
            shuffle=False)

        return validation_dataloader

    @torch.no_grad()
    def get_prompt_emb(self):

        if 'bert-' in self.args.backbone:
            prompt_emb = self.backbone.bert.embeddings.prompt_embeddings.weight

        if 'roberta-' in self.args.backbone:
            prompt_emb = self.backbone.roberta.embeddings.prompt_embeddings.weight

        if 't5-' in self.args.backbone:
            prompt_emb = self.backbone.encoder.embeddings.prompt_embeddings.weight

        return prompt_emb.detach().cpu()

    def train_prompt(self, model=None, task=None, **kwargs):
        self.args.output_dir = os.path.join(self.out_dir_root, 'prompt_emb')
        os.makedirs(self.args.output_dir, exist_ok=True)

        if model is None:
            model = self.model

        else:
            if isinstance(model, str):
                if model != self.args.backbone:
                    processor = data_processor_list[task]
                    model, template, verbalizer, plm, tokenizer, model_config, tokenizer_wrapper_class, model_type = self.get_model(model, processor, self.args)
                else:
                    model = self.model

            elif isinstance(model, torch.nn.Module):
                pass

        self._move_model_to_device(model, self.args.device)

        if task != self.args.dataset:
            processor = data_processor_list[self.args.dataset]()
            self.train_dataset = processor.train_dataset

        return super().train(**kwargs)

    def eval_prompt(self, model=None, eval_dataset=None, prompt_emb=None):
        r"""Evaluate a prompt"""

        self.args.output_dir = os.path.join(self.out_dir_root, 'prompt_emb')
        os.makedirs(self.args.output_dir, exist_ok=True)

        if model is not None and isinstance(model, str) and model != self.args.backbone:
            model, template, verbalizer, plm, tokenizer, model_config, tokenizer_wrapper_class, model_type = self.get_model(args.backbone, processor, args)
            self.model = model

        if prompt_emb is not None:
            if isinstance(prompt_emb, str):
                prompt_emb = torch.load(prompt_emb, map_location='cpu')
            elif isinstance(prompt_emb, torch.Tensor):
                pass
            else:
                raise NotImplementedError

            self.backbone.roberta.embeddings = nn.Parameter(prompt_emb.detach())

        if eval_dataset is None or eval_dataset == self.args.dataset:                    # Use default
            eval_dataset = self.eval_dataset

        elif eval_dataset != self.args.dataset:     # Use a dataset different from the source dataset
            processor = data_processor_list[eval_dataset]()
            eval_dataset = processor.eval_dataset

        else:
            raise NotImplementedError

        metrics = self.evaluate(eval_dataset=eval_dataset)
        self.save_metrics("eval_prompt", metrics)

        return metrics

    def cross_task_eval(self, model=None, target_task=None, prompt_emb=None):
        r"""Performs cross-task transfer evaluation."""

        self.args.output_dir = os.path.join(self.out_dir_root, 'prompt_emb')
        os.makedirs(self.args.output_dir, exist_ok=True)

        metrics = self.eval_prompt(model=model, eval_dataset=target_task, prompt_emb=prompt_emb)
        self.save_metrics("cross_task_eval", metrics)

        return metrics

    def cross_model_train(self, source_model=None, target_model=None, task=None, prompt_emb=None, **kwargs):
        r"""Performs corss-model transfer."""

        from .cross_model import CrossModelProjector

        self.args.output_dir = os.path.join(self.out_dir_root, 'cross_model_projector')
        os.makedirs(self.args.output_dir, exist_ok=True)

        processor = data_processor_list[task]()
        _, _, verbalizer, plm, tokenizer, target_model_config, tokenizer_wrapper_class, model_type = self.get_model(target_model, processor, self.args)
        template = CrossModelProjector(
            args=self.args, source_config=self.config, target_config=target_model_config,
            num_layers=self.args.num_proj_layers, flatten=self.args.flatten_proj,
            model=plm, text=self.template_text, tokenizer=tokenizer, num_tokens=self.args.prompt_len)
        self.model = PromptForClassification(plm=plm, template=template, verbalizer=verbalizer, freeze_plm=True)
        _keys_to_ignore_on_save = []
        for n, p in self.model.named_parameters():
            if not p.requires_grad:
                _keys_to_ignore_on_save.append(n)

        self.model._keys_to_ignore_on_save = _keys_to_ignore_on_save

        # Load source prompt or specific prompt
        # if prompt_emb is None:
        #     prompt_emb = os.path.join(self.out_dir_root, 'prompt_emb/checkpoint-711776/pytorch_model.bin')

        if isinstance(prompt_emb, str):
            prompt_emb = torch.load(prompt_emb, map_location='cpu')['prompt_model.template.soft_embeds'].to(self.args.device)
            self.model.prompt_model.template.soft_embeds = nn.Parameter(prompt_emb, requires_grad=False)
        elif isinstance(prompt_emb, torch.Tensor):
            self.model.prompt_model.template.soft_embeds = nn.Parameter(prompt_emb, requires_grad=False)

        self._move_model_to_device(self.model, self.args.device)

        trainable_parameter_names = []
        print('Trainable parameters:')
        for n, p in self.model.named_parameters():
            if p.requires_grad:
                print(n, p.shape)
                trainable_parameter_names.append(n)

        self.set_active_state_dict(self.model, trainable_parameter_names)

        if task != self.args.dataset:
            self.train_dataset = processor.train_dataset
            self.eval_dataset = processor.eval_dataset

        self.args.learning_rate = 0.1
        train_result = super().train(**kwargs)

        return train_result

    def cross_model_eval(self, source_model=None, target_model=None, task=None, prompt_emb=None):
        self.args.output_dir = os.path.join(self.out_dir_root, 'cross_model_projector')
        os.makedirs(self.args.output_dir, exist_ok=True)

        if task != self.args.dataset:
            processor = data_processor_list[task]()
            eval_dataset = processor.eval_dataset
        else:
            eval_dataset = self.eval_dataset

        metrics = self.evaluate(eval_dataset=eval_dataset)
        self.save_metrics("eval_prompt", metrics)

        return metrics

    def activated_neuron(self, model=None, task=None, layers=None, prompt_emb=None):
        r"""Get activated neuron."""

        processor = data_processor_list[task]()

        model, _, _, _, tokenizer, model_config, tokenizer_wrapper_class, model_type = self.get_model(model, processor, self.args)
        self._move_model_to_device(model, self.args.device)
        num_layers = model_config.num_hidden_layers

        if prompt_emb is not None:
            self.model.prompt_model.template.soft_embeds = nn.Parameter(prompt_emb, requires_grad=False)

        from openprompt.data_utils.utils import InputExample
        data = [InputExample(guid=0, text_a='<s>')]
        loader = PromptDataLoader(
            dataset=data,
            template=self.template,
            tokenizer=self.tokenizer,
            tokenizer_wrapper_class=self.tokenizer_wrapper_class,
            batch_size=1,
            max_seq_length=self.args.max_source_length,
            decoder_max_length=1,
            shuffle=False)

        outputs = [[] for _ in range(num_layers)]
        def save_ppt_outputs1_hook(n):
            def fn(_,__,output):
                outputs[n].append(output.detach().to('cpu'))
            return fn

        for n in range(num_layers):
            if model_type == 'bert' or model_type == 'roberta':
                model.prompt_model.plm.roberta.encoder.layer[n].intermediate.register_forward_hook(save_ppt_outputs1_hook(n))
            elif model_type == 't5':
                model.prompt_model.plm.t5.decoder.block[n].layer[2].DenseReluDense.wi.register_forward_hook(save_ppt_outputs1_hook(n))
            elif model_type == 'gpt':
                model.prompt_model.plm.gpt.transformer.h[n].mlp.c_fc.register_forward_hook(save_ppt_outputs1_hook(n))
            else:
                raise NotImplementedError

        for inputs in loader:
            _ = model(inputs.to(self.args.device))

        for k in range(num_layers):
            outputs[k] = torch.cat(outputs[k])

        outputs = torch.stack(outputs)
        outputs = outputs[:,:1,:1,:]    # Get the output of the <mask> position

        if layers is not None:
            outputs = outputs[torch.tensor(layers)]
            num_layers = len(layers)

        outputs = outputs.view(num_layers, -1)

        # Active neuron before ReLU
        torch.save(outputs, os.path.join(self.args.out_dir_root, self.dataset, 'activated_neuron_before_relu.pt'))

        # Active neuron after ReLU
        neuron_after_relu = (outputs > 0).int()
        torch.save(neuron_after_relu, os.path.join(self.args.out_dir_root, self.dataset, 'activated_neuron_after_relu.pt'))

        return outputs, neuron_after_relu

    def neuron_similarity(self, backbone=None, task1=None, task2=None):
        if task1 is None:
            task1 = self.args.source_dataset
        if task2 is None:
            task2 = self.args.target_dataset

        path1 = os.path.join(self.args.out_dir_root, self.args.source_dataset, 'activated_neuron_before_relu.pt')
        path2 = os.path.join(self.args.out_dir_root, self.args.target_dataset, 'activated_neuron_before_relu.pt')
        if os.path.exists(path1) == 0:
            raise ValueError(f"No neuron for {task1}.")
        if os.path.exists(path2) == 0:
            raise ValueError(f"No neuron for {task2}.")

        neuron1 = torch.load(path1)
        neuron2 = torch.load(path2)

        sim = F.cosine_similarity(neuron1, neuron2)

        return sim

    def mask_activated_neuron(self, model=None, task=None, layers=None, ratio=0.2):

        processor = data_processor_list[task]()

        model, _, _, _, tokenizer, model_config, tokenizer_wrapper_class, model_type = self.get_model(model, processor, self.args)
        self._move_model_to_device(model, self.args.device)
        num_layers = model_config.num_hidden_layers

        neuron = torch.load(f'outputs/{self.args.backbone}/{self.args.dataset}_{self.args.seed}/activated_neuron_before_relu.pt')
        original_shape = neuron.shape
        neuron = neuron.reshape(-1)
        mask = torch.ones_like(neuron)
        idx = torch.argsort(neuron, descending=True)
        idx = idx[:int(ratio * len(idx))]
        mask[idx] = 0
        mask = mask.view(original_shape)

        def save_ppt_outputs1_hook(n):
            def fn(_,__,output):
                output = output * mask[n].to('cuda')
                return output
            return fn

        if layers is None:
            layers = range(num_layers)

        for n in layers:
            if model_type == 'bert' or model_type == 'roberta':
                model.prompt_model.plm.roberta.encoder.layer[n].intermediate.register_forward_hook(save_ppt_outputs1_hook(n))
            elif model_type == 't5':
                model.prompt_model.plm.t5.decoder.block[n].layer[2].DenseReluDense.wi.register_forward_hook(save_ppt_outputs1_hook(n))
            elif model_type == 'gpt':
                model.prompt_model.plm.gpt.transformer.h[n].mlp.c_fc.register_forward_hook(save_ppt_outputs1_hook(n))
            else:
                raise NotImplementedError

        eval_results = self.eval_prompt(model=model, eval_dataset=task)

        return eval_results, mask

    def plot_neuron(self, model=None, task=None, **kwargs):
        import matplotlib.pyplot as plt
        import seaborn as sns

        plt.figure(figsize=(40, 20))
        neuron = torch.load(f'outputs/{self.args.backbone}/{self.args.dataset}_{self.args.seed}/activated_neuron_before_relu.pt')
        sns.heatmap(neuron.numpy(), cmap='Reds', **kwargs)
        plt.xlabel('Neuron')
        plt.ylabel('Layer')

    def set_active_state_dict(self, module: nn.Module, includes=['prompt_model.template.soft_embeds']):
        r"""modify the state_dict function of the model (by default, the backbone model) to return only the tunable part.
        Args:
            module (:obj:`nn.Module`): The module modified. The modification is in-place.
        """

        def _caller(_org_func, includes,  *args, **kwargs):
            state_dict = _org_func(*args, **kwargs)
            keys = list(state_dict.keys())
            for n  in keys:
                if n not in includes:
                    state_dict.pop(n)
            return state_dict
        includes = includes          # use excludes will have trouble when the model have shared weights
        if hasattr(module.state_dict, "__wrapped__"):
            raise RuntimeWarning("The forward function might have been wrapped by a decorator, is it intended? Do you freeze the parameters twice?")
        module.state_dict = decorate(module.state_dict, _caller, extras=(includes,), kwsyntax=True)