import hashlib
import hmac
import json
import os
import posixpath
import re

from six.moves import http_client
from six.moves import urllib
from six.moves.urllib.parse import urljoin

from google.auth import _helpers
from google.auth import environment_vars
from google.auth import exceptions
from google.auth import external_account

# AWS Signature Version 4 signing algorithm identifier.
_AWS_ALGORITHM = "AWS4-HMAC-SHA256"
# The termination string for the AWS credential scope value as defined in
# https://docs.aws.amazon.com/general/latest/gr/sigv4-create-string-to-sign.html
_AWS_REQUEST_TYPE = "aws4_request"
# The AWS authorization header name for the security session token if available.
_AWS_SECURITY_TOKEN_HEADER = "x-amz-security-token"
# The AWS authorization header name for the auto-generated date.
_AWS_DATE_HEADER = "x-amz-date"

class RequestSigner(object):
    """Implements an AWS request signer based on the AWS Signature Version 4 signing
    process.
    https://docs.aws.amazon.com/general/latest/gr/signature-version-4.html
    """

    def __init__(self, region_name):
        """Instantiates an AWS request signer used to compute authenticated signed
        requests to AWS APIs based on the AWS Signature Version 4 signing process.

        Args:
            region_name (str): The AWS region to use.
        """

        self._region_name = region_name

    def get_request_options(
        self,
        aws_security_credentials,
        url,
        method,
        request_payload="",
        additional_headers={},
    ):
        """Generates the signed request for the provided HTTP request for calling
        an AWS API. This follows the steps described at:
        https://docs.aws.amazon.com/general/latest/gr/sigv4_signing.html

        Args:
            aws_security_credentials (Mapping[str, str]): A dictionary containing
                the AWS security credentials.
            url (str): The AWS service URL containing the canonical URI and
                query string.
            method (str): The HTTP method used to call this API.
            request_payload (Optional[str]): The optional request payload if
                available.
            additional_headers (Optional[Mapping[str, str]]): The optional
                additional headers needed for the requested AWS API.

        Returns:
            Mapping[str, str]: The AWS signed request dictionary object.
        """
        # Get AWS credentials.
        access_key = aws_security_credentials.get("access_key_id")
        secret_key = aws_security_credentials.get("secret_access_key")
        security_token = aws_security_credentials.get("security_token")

        additional_headers = additional_headers or {}

        uri = urllib.parse.urlparse(url)
        # Normalize the URL path. This is needed for the canonical_uri.
        # os.path.normpath can't be used since it normalizes "/" paths
        # to "\\" in Windows OS.
        normalized_uri = urllib.parse.urlparse(
            urljoin(url, posixpath.normpath(uri.path))
        )
        # Validate provided URL.
        if not uri.hostname or uri.scheme != "https":
            raise exceptions.InvalidResource("Invalid AWS service URL")

        header_map = _generate_authentication_header_map(
            host=uri.hostname,
            canonical_uri=normalized_uri.path or "/",
            canonical_querystring=_get_canonical_querystring(uri.query),
            method=method,
            region=self._region_name,
            access_key=access_key,
            secret_key=secret_key,
            security_token=security_token,
            request_payload=request_payload,
            additional_headers=additional_headers,
        )
        headers = {
            "Authorization": header_map.get("authorization_header"),
            "host": uri.hostname,
        }
        # Add x-amz-date if available.
        if "amz_date" in header_map:
            headers[_AWS_DATE_HEADER] = header_map.get("amz_date")
        # Append additional optional headers, eg. X-Amz-Target, Content-Type, etc.
        for key in additional_headers:
            headers[key] = additional_headers[key]

        # Add session token if available.
        if security_token is not None:
            headers[_AWS_SECURITY_TOKEN_HEADER] = security_token

        signed_request = {"url": url, "method": method, "headers": headers}
        if request_payload:
            signed_request["data"] = request_payload
        return signed_request

def _get_canonical_querystring(query):
    """Generates the canonical query string given a raw query string.
    Logic is based on
    https://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html

    Args:
        query (str): The raw query string.

    Returns:
        str: The canonical query string.
    """
    # Parse raw query string.
    querystring = urllib.parse.parse_qs(query)
    querystring_encoded_map = {}
    for key in querystring:
        quote_key = urllib.parse.quote(key, safe="-_.~")
        # URI encode key.
        querystring_encoded_map[quote_key] = []
        for item in querystring[key]:
            # For each key, URI encode all values for that key.
            querystring_encoded_map[quote_key].append(
                urllib.parse.quote(item, safe="-_.~")
            )
        # Sort values for each key.
        querystring_encoded_map[quote_key].sort()
    # Sort keys.
    sorted_keys = list(querystring_encoded_map.keys())
    sorted_keys.sort()
    # Reconstruct the query string. Preserve keys with multiple values.
    querystring_encoded_pairs = []
    for key in sorted_keys:
        for item in querystring_encoded_map[key]:
            querystring_encoded_pairs.append("{}={}".format(key, item))
    return "&".join(querystring_encoded_pairs)

def _sign(key, msg):
    """Creates the HMAC-SHA256 hash of the provided message using the provided
    key.

    Args:
        key (str): The HMAC-SHA256 key to use.
        msg (str): The message to hash.

    Returns:
        str: The computed hash bytes.
    """
    return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()

def _get_signing_key(key, date_stamp, region_name, service_name):
    """Calculates the signing key used to calculate the signature for
    AWS Signature Version 4 based on:
    https://docs.aws.amazon.com/general/latest/gr/sigv4-calculate-signature.html

    Args:
        key (str): The AWS secret access key.
        date_stamp (str): The '%Y%m%d' date format.
        region_name (str): The AWS region.
        service_name (str): The AWS service name, eg. sts.

    Returns:
        str: The signing key bytes.
    """
    k_date = _sign(("AWS4" + key).encode("utf-8"), date_stamp)
    k_region = _sign(k_date, region_name)
    k_service = _sign(k_region, service_name)
    k_signing = _sign(k_service, "aws4_request")
    return k_signing

def _generate_authentication_header_map(
    host,
    canonical_uri,
    canonical_querystring,
    method,
    region,
    access_key,
    secret_key,
    security_token,
    request_payload="",
    additional_headers={},
):
    """Generates the authentication header map needed for generating the AWS
    Signature Version 4 signed request.

    Args:
        host (str): The AWS service URL hostname.
        canonical_uri (str): The AWS service URL path name.
        canonical_querystring (str): The AWS service URL query string.
        method (str): The HTTP method used to call this API.
        region (str): The AWS region.
        access_key (str): The AWS access key ID.
        secret_key (str): The AWS secret access key.
        security_token (Optional[str]): The AWS security session token. This is
            available for temporary sessions.
        request_payload (Optional[str]): The optional request payload if
            available.
        additional_headers (Optional[Mapping[str, str]]): The optional
            additional headers needed for the requested AWS API.

    Returns:
        Mapping[str, str]: The AWS authentication header dictionary object.
            This contains the x-amz-date and authorization header information.
    """
    # iam.amazonaws.com host => iam service.
    # sts.us-east-2.amazonaws.com host => sts service.
    service_name = host.split(".")[0]

    current_time = _helpers.utcnow()
    amz_date = current_time.strftime("%Y%m%dT%H%M%SZ")
    date_stamp = current_time.strftime("%Y%m%d")

    # Change all additional headers to be lower case.
    full_headers = {}
    for key in additional_headers:
        full_headers[key.lower()] = additional_headers[key]
    # Add AWS session token if available.
    if security_token is not None:
        full_headers[_AWS_SECURITY_TOKEN_HEADER] = security_token

    # Required headers
    full_headers["host"] = host
    # Do not use generated x-amz-date if the date header is provided.
    # Previously the date was not fixed with x-amz- and could be provided
    # manually.
    # https://github.com/boto/botocore/blob/879f8440a4e9ace5d3cf145ce8b3d5e5ffb892ef/tests/unit/auth/aws4_testsuite/get-header-value-trim.req
    if "date" not in full_headers:
        full_headers[_AWS_DATE_HEADER] = amz_date

    # Header keys need to be sorted alphabetically.
    canonical_headers = ""
    header_keys = list(full_headers.keys())
    header_keys.sort()
    for key in header_keys:
        canonical_headers = "{}{}:{}\n".format(
            canonical_headers, key, full_headers[key]
        )
    signed_headers = ";".join(header_keys)

    payload_hash = hashlib.sha256((request_payload or "").encode("utf-8")).hexdigest()

    # https://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html
    canonical_request = "{}\n{}\n{}\n{}\n{}\n{}".format(
        method,
        canonical_uri,
        canonical_querystring,
        canonical_headers,
        signed_headers,
        payload_hash,
    )

    credential_scope = "{}/{}/{}/{}".format(
        date_stamp, region, service_name, _AWS_REQUEST_TYPE
    )

    # https://docs.aws.amazon.com/general/latest/gr/sigv4-create-string-to-sign.html
    string_to_sign = "{}\n{}\n{}\n{}".format(
        _AWS_ALGORITHM,
        amz_date,
        credential_scope,
        hashlib.sha256(canonical_request.encode("utf-8")).hexdigest(),
    )

    # https://docs.aws.amazon.com/general/latest/gr/sigv4-calculate-signature.html
    signing_key = _get_signing_key(secret_key, date_stamp, region, service_name)
    signature = hmac.new(
        signing_key, string_to_sign.encode("utf-8"), hashlib.sha256
    ).hexdigest()

    # https://docs.aws.amazon.com/general/latest/gr/sigv4-add-signature-to-request.html
    authorization_header = "{} Credential={}/{}, SignedHeaders={}, Signature={}".format(
        _AWS_ALGORITHM, access_key, credential_scope, signed_headers, signature
    )

    authentication_header = {"authorization_header": authorization_header}
    # Do not use generated x-amz-date if the date header is provided.
    if "date" not in full_headers:
        authentication_header["amz_date"] = amz_date
    return authentication_header

class Credentials(external_account.Credentials):
    """AWS external account credentials.
    This is used to exchange serialized AWS signature v4 signed requests to
    AWS STS GetCallerIdentity service for Google access tokens.
    """

    def __init__(
        self,
        audience,
        subject_token_type,
        token_url,
        credential_source=None,
        *args,
        **kwargs
    ):
        """Instantiates an AWS workload external account credentials object.

        Args:
            audience (str): The STS audience field.
            subject_token_type (str): The subject token type.
            token_url (str): The STS endpoint URL.
            credential_source (Mapping): The credential source dictionary used
                to provide instructions on how to retrieve external credential
                to be exchanged for Google access tokens.
            args (List): Optional positional arguments passed into the underlying :meth:`~external_account.Credentials.__init__` method.
            kwargs (Mapping): Optional keyword arguments passed into the underlying :meth:`~external_account.Credentials.__init__` method.

        Raises:
            google.auth.exceptions.RefreshError: If an error is encountered during
                access token retrieval logic.
            ValueError: For invalid parameters.

        .. note:: Typically one of the helper constructors
            :meth:`from_file` or
            :meth:`from_info` are used instead of calling the constructor directly.
        """
        super(Credentials, self).__init__(
            audience=audience,
            subject_token_type=subject_token_type,
            token_url=token_url,
            credential_source=credential_source,
            *args,
            **kwargs
        )
        credential_source = credential_source or {}
        self._environment_id = credential_source.get("environment_id") or ""
        self._region_url = credential_source.get("region_url")
        self._security_credentials_url = credential_source.get("url")
        self._cred_verification_url = credential_source.get(
            "regional_cred_verification_url"
        )
        self._imdsv2_session_token_url = credential_source.get(
            "imdsv2_session_token_url"
        )
        self._region = None
        self._request_signer = None
        self._target_resource = audience

        # Get the environment ID. Currently, only one version supported (v1).
        matches = re.match(r"^(aws)([\d]+)$", self._environment_id)
        if matches:
            env_id, env_version = matches.groups()
        else:
            env_id, env_version = (None, None)

        if env_id != "aws" or self._cred_verification_url is None:
            raise exceptions.InvalidResource(
                "No valid AWS 'credential_source' provided"
            )
        elif int(env_version or "") != 1:
            raise exceptions.InvalidValue(
                "aws version '{}' is not supported in the current build.".format(
                    env_version
                )
            )

    def retrieve_subject_token(self, request):
        """Retrieves the subject token using the credential_source object.
        The subject token is a serialized `AWS GetCallerIdentity signed request`_.

        The logic is summarized as:

        Retrieve the AWS region from the AWS_REGION or AWS_DEFAULT_REGION
        environment variable or from the AWS metadata server availability-zone
        if not found in the environment variable.

        Check AWS credentials in environment variables. If not found, retrieve
        from the AWS metadata server security-credentials endpoint.

        When retrieving AWS credentials from the metadata server
        security-credentials endpoint, the AWS role needs to be determined by
        calling the security-credentials endpoint without any argument. Then the
        credentials can be retrieved via: security-credentials/role_name

        Generate the signed request to AWS STS GetCallerIdentity action.

        Inject x-goog-cloud-target-resource into header and serialize the
        signed request. This will be the subject-token to pass to GCP STS.

        .. _AWS GetCallerIdentity signed request:
            https://cloud.google.com/iam/docs/access-resources-aws#exchange-token

        Args:
            request (google.auth.transport.Request): A callable used to make
                HTTP requests.
        Returns:
            str: The retrieved subject token.
        """
        # Fetch the session token required to make meta data endpoint calls to aws.
        if (
            request is not None
            and self._imdsv2_session_token_url is not None
            and self._should_use_metadata_server()
        ):
            headers = {"X-aws-ec2-metadata-token-ttl-seconds": "300"}

            imdsv2_session_token_response = request(
                url=self._imdsv2_session_token_url, method="PUT", headers=headers
            )

            if imdsv2_session_token_response.status != 200:
                raise exceptions.RefreshError(
                    "Unable to retrieve AWS Session Token",
                    imdsv2_session_token_response.data,
                )

            imdsv2_session_token = imdsv2_session_token_response.data
        else:
            imdsv2_session_token = None

        # Initialize the request signer if not yet initialized after determining
        # the current AWS region.
        if self._request_signer is None:
            self._region = self._get_region(
                request, self._region_url, imdsv2_session_token
            )
            self._request_signer = RequestSigner(self._region)

        # Retrieve the AWS security credentials needed to generate the signed
        # request.
        aws_security_credentials = self._get_security_credentials(
            request, imdsv2_session_token
        )
        # Generate the signed request to AWS STS GetCallerIdentity API.
        # Use the required regional endpoint. Otherwise, the request will fail.
        request_options = self._request_signer.get_request_options(
            aws_security_credentials,
            self._cred_verification_url.replace("{region}", self._region),
            "POST",
        )
        # The GCP STS endpoint expects the headers to be formatted as:
        # [
        #   {key: 'x-amz-date', value: '...'},
        #   {key: 'Authorization', value: '...'},
        #   ...
        # ]
        # And then serialized as:
        # quote(json.dumps({
        #   url: '...',
        #   method: 'POST',
        #   headers: [{key: 'x-amz-date', value: '...'}, ...]
        # }))
        request_headers = request_options.get("headers")
        # The full, canonical resource name of the workload identity pool
        # provider, with or without the HTTPS prefix.
        # Including this header as part of the signature is recommended to
        # ensure data integrity.
        request_headers["x-goog-cloud-target-resource"] = self._target_resource

        # Serialize AWS signed request.
        # Keeping inner keys in sorted order makes testing easier for Python
        # versions <=3.5 as the stringified JSON string would have a predictable
        # key order.
        aws_signed_req = {}
        aws_signed_req["url"] = request_options.get("url")
        aws_signed_req["method"] = request_options.get("method")
        aws_signed_req["headers"] = []
        # Reformat header to GCP STS expected format.
        for key in sorted(request_headers.keys()):
            aws_signed_req["headers"].append(
                {"key": key, "value": request_headers[key]}
            )

        return urllib.parse.quote(
            json.dumps(aws_signed_req, separators=(",", ":"), sort_keys=True)
        )

    def _get_region(self, request, url, imdsv2_session_token):
        """Retrieves the current AWS region from either the AWS_REGION or
        AWS_DEFAULT_REGION environment variable or from the AWS metadata server.

        Args:
            request (google.auth.transport.Request): A callable used to make
                HTTP requests.
            url (str): The AWS metadata server region URL.
            imdsv2_session_token (str): The AWS IMDSv2 session token to be added as a
                header in the requests to AWS metadata endpoint.

        Returns:
            str: The current AWS region.

        Raises:
            google.auth.exceptions.RefreshError: If an error occurs while
                retrieving the AWS region.
        """
        # The AWS metadata server is not available in some AWS environments
        # such as AWS lambda. Instead, it is available via environment
        # variable.
        env_aws_region = os.environ.get(environment_vars.AWS_REGION)
        if env_aws_region is not None:
            return env_aws_region

        env_aws_region = os.environ.get(environment_vars.AWS_DEFAULT_REGION)
        if env_aws_region is not None:
            return env_aws_region

        if not self._region_url:
            raise exceptions.RefreshError("Unable to determine AWS region")

        headers = None
        if imdsv2_session_token is not None:
            headers = {"X-aws-ec2-metadata-token": imdsv2_session_token}

        response = request(url=self._region_url, method="GET", headers=headers)

        # Support both string and bytes type response.data.
        response_body = (
            response.data.decode("utf-8")
            if hasattr(response.data, "decode")
            else response.data
        )

        if response.status != 200:
            raise exceptions.RefreshError(
                "Unable to retrieve AWS region", response_body
            )

        # This endpoint will return the region in format: us-east-2b.
        # Only the us-east-2 part should be used.
        return response_body[:-1]

    def _get_security_credentials(self, request, imdsv2_session_token):
        """Retrieves the AWS security credentials required for signing AWS
        requests from either the AWS security credentials environment variables
        or from the AWS metadata server.

        Args:
            request (google.auth.transport.Request): A callable used to make
                HTTP requests.
            imdsv2_session_token (str): The AWS IMDSv2 session token to be added as a
                header in the requests to AWS metadata endpoint.

        Returns:
            Mapping[str, str]: The AWS security credentials dictionary object.

        Raises:
            google.auth.exceptions.RefreshError: If an error occurs while
                retrieving the AWS security credentials.
        """

        # Check environment variables for permanent credentials first.
        # https://docs.aws.amazon.com/general/latest/gr/aws-sec-cred-types.html
        env_aws_access_key_id = os.environ.get(environment_vars.AWS_ACCESS_KEY_ID)
        env_aws_secret_access_key = os.environ.get(
            environment_vars.AWS_SECRET_ACCESS_KEY
        )
        # This is normally not available for permanent credentials.
        env_aws_session_token = os.environ.get(environment_vars.AWS_SESSION_TOKEN)
        if env_aws_access_key_id and env_aws_secret_access_key:
            return {
                "access_key_id": env_aws_access_key_id,
                "secret_access_key": env_aws_secret_access_key,
                "security_token": env_aws_session_token,
            }

        # Get role name.
        role_name = self._get_metadata_role_name(request, imdsv2_session_token)

        # Get security credentials.
        credentials = self._get_metadata_security_credentials(
            request, role_name, imdsv2_session_token
        )

        return {
            "access_key_id": credentials.get("AccessKeyId"),
            "secret_access_key": credentials.get("SecretAccessKey"),
            "security_token": credentials.get("Token"),
        }

    def _get_metadata_security_credentials(
        self, request, role_name, imdsv2_session_token
    ):
        """Retrieves the AWS security credentials required for signing AWS
        requests from the AWS metadata server.

        Args:
            request (google.auth.transport.Request): A callable used to make
                HTTP requests.
            role_name (str): The AWS role name required by the AWS metadata
                server security_credentials endpoint in order to return the
                credentials.
            imdsv2_session_token (str): The AWS IMDSv2 session token to be added as a
                header in the requests to AWS metadata endpoint.

        Returns:
            Mapping[str, str]: The AWS metadata server security credentials
                response.

        Raises:
            google.auth.exceptions.RefreshError: If an error occurs while
                retrieving the AWS security credentials.
        """
        headers = {"Content-Type": "application/json"}
        if imdsv2_session_token is not None:
            headers["X-aws-ec2-metadata-token"] = imdsv2_session_token

        response = request(
            url="{}/{}".format(self._security_credentials_url, role_name),
            method="GET",
            headers=headers,
        )

        # support both string and bytes type response.data
        response_body = (
            response.data.decode("utf-8")
            if hasattr(response.data, "decode")
            else response.data
        )

        if response.status != http_client.OK:
            raise exceptions.RefreshError(
                "Unable to retrieve AWS security credentials", response_body
            )

        credentials_response = json.loads(response_body)

        return credentials_response

    def _get_metadata_role_name(self, request, imdsv2_session_token):
        """Retrieves the AWS role currently attached to the current AWS
        workload by querying the AWS metadata server. This is needed for the
        AWS metadata server security credentials endpoint in order to retrieve
        the AWS security credentials needed to sign requests to AWS APIs.

        Args:
            request (google.auth.transport.Request): A callable used to make
                HTTP requests.
            imdsv2_session_token (str): The AWS IMDSv2 session token to be added as a
                header in the requests to AWS metadata endpoint.

        Returns:
            str: The AWS role name.

        Raises:
            google.auth.exceptions.RefreshError: If an error occurs while
                retrieving the AWS role name.
        """
        if self._security_credentials_url is None:
            raise exceptions.RefreshError(
                "Unable to determine the AWS metadata server security credentials endpoint"
            )

        headers = None
        if imdsv2_session_token is not None:
            headers = {"X-aws-ec2-metadata-token": imdsv2_session_token}

        response = request(
            url=self._security_credentials_url, method="GET", headers=headers
        )

        # support both string and bytes type response.data
        response_body = (
            response.data.decode("utf-8")
            if hasattr(response.data, "decode")
            else response.data
        )

        if response.status != http_client.OK:
            raise exceptions.RefreshError(
                "Unable to retrieve AWS role name", response_body
            )

        return response_body

    def _should_use_metadata_server(self):
        # The AWS region can be provided through AWS_REGION or AWS_DEFAULT_REGION.
        # The metadata server should be used if it cannot be retrieved from one of
        # these environment variables.
        if not os.environ.get(environment_vars.AWS_REGION) and not os.environ.get(
            environment_vars.AWS_DEFAULT_REGION
        ):
            return True

        # AWS security credentials can be retrieved from the AWS_ACCESS_KEY_ID
        # and AWS_SECRET_ACCESS_KEY environment variables. The metadata server
        # should be used if either of these are not available.
        if not os.environ.get(environment_vars.AWS_ACCESS_KEY_ID) or not os.environ.get(
            environment_vars.AWS_SECRET_ACCESS_KEY
        ):
            return True

        return False

    @classmethod
    def from_info(cls, info, **kwargs):
        """Creates an AWS Credentials instance from parsed external account info.

        Args:
            info (Mapping[str, str]): The AWS external account info in Google
                format.
            kwargs: Additional arguments to pass to the constructor.

        Returns:
            google.auth.aws.Credentials: The constructed credentials.

        Raises:
            ValueError: For invalid parameters.
        """
        return super(Credentials, cls).from_info(info, **kwargs)

    @classmethod
    def from_file(cls, filename, **kwargs):
        """Creates an AWS Credentials instance from an external account json file.

        Args:
            filename (str): The path to the AWS external account json file.
            kwargs: Additional arguments to pass to the constructor.

        Returns:
            google.auth.aws.Credentials: The constructed credentials.
        """
        return super(Credentials, cls).from_file(filename, **kwargs)


import json

#filename = 'log_1_200bx_braketest_noabs_nopulse.json'
filename = 'log_2_200bx_braketest_abs_nopulse.json'

file = open(filename, 'r')
data = json.loads(file.read())
params = open('params', 'r').read().splitlines()
params.remove(params[-1])

results = []
n = 0
for i in data['electrics']['rpm']['values']:
    results.append({})
    if type(data['electrics']['rpm']['values'][n]) != dict:
        n += 1
        continue
    results[n]['time'] = data['electrics']['rpm']['values'][n]['time']
    for param in params:
            results[n][param] = data['electrics'][param]['values'][n]['value']
    n += 1

last_time = None
last_odometer_airspeed_mps = 0

params.insert(params.index('wheelspeed'), 'airspeed_mph')
params.insert(params.index('rpm'), 'wheelspeed_mph')
params.insert(params.index('airspeed'), 'odometer_airspeed_meters')

# calculate other parameters not originally included in the data
for i in results:
    try:
        # convert meters per second to miles per hour (conversion factor 2.237)
        i['airspeed_mph'] = 2.237*i['airspeed']
        i['wheelspeed_mph'] = 2.237*i['wheelspeed']
        # calculate how far we've traveled in the air.
        # (regular odometer param (also in meters) is based on how far the wheels have turned, like a real car)
        if last_time != None:
            time_diff = (i['time'] - last_time) / 1000
            i['odometer_airspeed_meters'] = last_odometer_airspeed_mps + i['airspeed'] * time_diff
            last_odometer_airspeed_mps = i['odometer_airspeed_meters']
        last_time = i['time']
    except KeyError as e:
        #print(e)
        #print(results.index(i))
        pass

# jankily construct a csv file
# first the names of the columns
s = 'time,'
for p in params:
    s += p + ','
s += '\n'
# then the actual values
for i in results:
    s2 = ''
    try:
        s2 += str(i['time'])
        s2 += ','
        for p in params:
            s2 += str(i[p])
            s2 += ','
        s2 += '\n'
    except KeyError as e:
        #print(e)
        continue
    s += s2

outfile = open(filename + '.csv', 'w')
outfile.write(s)

# grab braking distances
last_airspeed = None
for i in results:
    try:
        if last_airspeed != None and last_airspeed > 60 and i['airspeed_mph'] < 60:
            s = ''
            s += f"{ i['odometer_airspeed_meters'] }, { i['time'] }, "
        if last_airspeed != None and last_airspeed > 1 and i['airspeed_mph'] < 1:
            s += f"{ i['odometer_airspeed_meters'] }, { i['time'] } "
            print(s)
        last_airspeed = i['airspeed_mph']
    except KeyError as e:
        #print(e)
        #print(results.index(i))
        continue


import datetime
import matplotlib.pyplot as plt
import pandas.io.sql as sqlio
from nyc_taxi_trips_load_tasks.common.aws import upload_file_to_s3
import pandas as pd
from airflow.contrib.hooks.aws_hook import AwsHook
from airflow.hooks.postgres_hook import PostgresHook
import os
import logging

def build_query_for(year, month):
    return """
        SELECT
            CAST(SPLIT_PART(REPLACE(amount_group, '$', ''), '-', 1) AS NUMERIC) AS amount_group_start_range,
            amount_group,
            frequency
        FROM
            bi_monthly_price_frequency
        WHERE
            pickup_month = '{}-{}'
        ORDER BY 1
    """.format(year, '0' + str(month) if month < 10 else month)

def build_graphic_for(execution_date, conn, credentials):
    sql_query = build_query_for(execution_date.year, execution_date.month)
    logging.info(sql_query)
    df = sqlio.read_sql_query(sql_query, conn)

    if df.empty:
        logging.info('Query returned empty results set')
        return

    df.plot(kind='bar', x='amount_group', y='frequency')
    img_name = "{}-{}_monthly_price_frequency.png".format(execution_date.year, execution_date.month)
    plt.savefig(img_name, bbox_inches='tight')

    upload_file_to_s3(img_name, "data-sprints-eng-test/outputs/monthly/{}".format(img_name), credentials)

    os.remove(img_name)

def generate_monthly_graphics(*args, **kwargs):
    aws_hook = AwsHook("aws_credentials")
    credentials = aws_hook.get_credentials()
    redshift_hook = PostgresHook("redshift")
    conn = redshift_hook.get_conn()

    # # #
    execution_date = datetime.datetime.strptime(kwargs["ds"], '%Y-%m-%d')
    # # #

    build_graphic_for(execution_date, conn, credentials)


def grow(arr):
    x = 1
    for i in arr:
        x = i * x
    return x

# P int
# R int with each number in the list multiplied with the next one
# E [2,3,4] -> 2 * 3 * 4 = 24
# P iterate the list and multiply with the previous number, stored in x

10.16.23 multiply numbers in order in a list.py


r"""
    Module associated with booleval, logictable, and logicparser, to do
    boolean evaluation of boolean formulas.
    Formulas consist of the operators &, |, ~, ^, ->, <->, corresponding
    to and, or, not, xor, if...then, if and only if.  Operators can
    be applied to variables that consist of a leading letter and trailing
    underscores and alphanumerics.  Parentheses may be used to
    explicitly show order of operation.

    AUTHORS:
        -- Chris Gorecki

    EXAMPLES:
        sage: import sage.logic.propcalc as propcalc
        sage: f = propcalc.formula("a&((b|c)^a->c)<->b")
        sage: g = propcalc.formula("boolean<->algebra")
        sage: (f&~g).ifthen(f)
        ((a&((b|c)^a->c)<->b)&(~(boolean<->algebra)))->(a&((b|c)^a->c)<->b)

     EXAMPLES:
        sage: import sage.logic.propcalc as propcalc
        sage: f = propcalc.formula("a&((b|c)^a->c)<->b")
        sage: g = propcalc.formula("boolean<->algebra")
        sage: (f&~g).ifthen(f)
        ((a&((b|c)^a->c)<->b)&(~(boolean<->algebra)))->(a&((b|c)^a->c)<->b)

    We can create a truth table from a formula.
        sage: f.truthtable()
        a      b      c      value
        False  False  False  True
        False  False  True   True
        False  True   False  False
        False  True   True   False
        True   False  False  True
        True   False  True   False
        True   True   False  True
        True   True   True   True
        sage: f.truthtable(end=3)
        a      b      c      value
        False  False  False  True
        False  False  True   True
        False  True   False  False
        sage: f.truthtable(start=4)
        a      b      c      value
        True   False  False  True
        True   False  True   False
        True   True   False  True
        True   True   True   True
        sage: propcalc.formula("a").truthtable()
        a      value
        False  False
        True   True

    Now we can evaluate the formula for a given set of inputs.
        sage: f.evaluate({'a':True, 'b':False, 'c':True})
        False
        sage: f.evaluate({'a':False, 'b':False, 'c':True})
        True

    And we can convert a boolean formula to conjunctive normal form.
        sage: f.convert_cnf_table()
        sage: f
        (a|~b|c)&(a|~b|~c)&(~a|b|~c)
        sage: f.convert_cnf_recur()
        sage: f
        (a|~b|c)&(a|~b|~c)&(~a|b|~c)

    Or determine if an expression is satisfiable, a contradiction, or a tautology.
        sage: f = propcalc.formula("a|b")
        sage: f.is_satisfiable()
        True
        sage: f = f & ~f
        sage: f.is_satisfiable()
        False
        sage: f.is_contradiction()
        True
        sage: f = f | ~f
        sage: f.is_tautology()
        True

    The equality operator compares semantic equivalence.
        sage: f = propcalc.formula("(a|b)&c")
        sage: g = propcalc.formula("c&(b|a)")
        sage: f == g
        True
        sage: g = propcalc.formula("a|b&c")
        sage: f == g
        False

    It is an error to create a formula with bad syntax.
        sage: propcalc.formula("")
        Traceback (most recent call last):
        ...
        SyntaxError: malformed statement
        sage: propcalc.formula("a&b~(c|(d)")
        Traceback (most recent call last):
        ...
        SyntaxError: malformed statement
        sage: propcalc.formula("a&&b")
        Traceback (most recent call last):
        ...
        SyntaxError: malformed statement
        sage: propcalc.formula("a&b a")
        Traceback (most recent call last):
        ...
        SyntaxError: malformed statement

        It is also an error to not abide by the naming conventions.
        sage: propcalc.formula("~a&9b")
        Traceback (most recent call last):
        ...
        NameError: invalid variable name 9b: identifiers must begin with a letter and contain only alphanumerics and underscores

"""
#*******************************************************************************************
#copyright (C) 2006 William Stein <wstein@gmail.com>
#copyright (C) 2006 Chris Gorecki <chris.k.gorecki@gmail.com>
#Distributed under the terms of the GNU General Public License (GPL)
#http://www.gnu.org/licenses/
#*******************************************************************************************
import booleval
import logictable
import logicparser
#import boolopt
from types import *

latex_operators = [('&', '\\wedge '),
                   ('|', '\\vee '),
                   ('~', '\\neg '),
                   ('^', '\\oplus '),
                   ('<->', '\\leftrightarrow '),
                   ('->', '\\rightarrow ')]

class BooleanFormula:
    __expression = ""
    __tree = []
    __vars_order = []

    def __init__(self, exp, tree, vo):
        r"""
        This function initializes the data fields and is called when a
        new statement is created.

        INPUT:
            self -- the calling object.
            exp -- a string containing the logic expression to be manipulated.
            tree -- a list containing the parse tree of the expression.
            vo -- a list of the variables in the expression in order,
                  with each variable occurring only once.

        OUTPUT:
            Effectively returns an instance of this class.

        EXAMPLES:
        This example illustrates the creation of a statement.
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b|~(c|a)")
            sage: s
            a&b|~(c|a)
        """
        self.__expression = ""
        #remove white space from expression
        for c in exp:
            if(c != ' '):
                self.__expression += c
        self.__tree = tree
        self.__vars_order = vo

    def __repr__(self):
        r"""
        Returns a string representation of this statement.

        INPUT:
            self -- the calling object.

        OUTPUT:
            Returns the string representation of this statement.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: propcalc.formula("man->monkey&human")
            man->monkey&human
        """
        return self.__expression

    def _latex_(self):
        r"""
        Returns a LaTeX  representation of this statement.

        INPUT:
            self -- the calling object.

        OUTPUT:
            Returns the latex representation of this statement.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("man->monkey&human")
            sage: latex(s)
            man\rightarrow monkey\wedge human

            sage: f = propcalc.formula("a & ((~b | c) ^ a -> c) <-> ~b")
            sage: latex(f)
            a\wedge ((\neg b\vee c)\oplus a\rightarrow c)\leftrightarrow \neg b
        """
        latex_expression = self.__expression
        for old, new in latex_operators:
            latex_expression = latex_expression.replace(old, new)
        return latex_expression

    def tree(self):
        r"""
        Returns a list containing the parse tree of this boolean expression.

        INPUT:
            self -- the calling object.

        OUTPUT:
            A list containing the parse tree of self.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("man -> monkey & human")
            sage: s.tree()
            ['->', 'man', ['&', 'monkey', 'human']]
            sage: f = propcalc.formula("a & ((~b | c) ^ a -> c) <-> ~b")
            sage: f.tree()
            <BLANKLINE>
            ['<->',
             ['&', 'a', ['->', ['^', ['|', ['~', 'b', None], 'c'], 'a'], 'c']],
             ['~', 'b', None]]
        """
        return self.__tree

    def __or__(self, other):
        r"""
        Overloads the | operator to 'or' two statements together.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns a new statement that is the first statement logically
                or'ed together.
        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: f = propcalc.formula("c^d")
            sage: s | f
            (a&b)|(c^d)
        """
        return self.add_statement(other, '|')

    def __and__(self, other):
        r"""
        Overloads the & operator to 'and' two statements together.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns a new statement that is the first statement logically
            and'ed together.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: f = propcalc.formula("c^d")
            sage: s & f
            (a&b)&(c^d)
        """
        return self.add_statement(other, '&')

    def __xor__(self, other):
        r"""
        Overloads the ^ operator to xor two statements together.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns a new statement that is the first statement logically
            xor'ed together.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: f = propcalc.formula("c^d")
            sage: s ^ f
            (a&b)^(c^d)
        """
        return self.add_statement(other, '^')

    def __pow__(self, other):
        r"""
        Overloads the ^ operator to xor two statements together.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns a new statement that is the first statement logically
            xor'ed together.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: f = propcalc.formula("c^d")
            sage: s ^ f
            (a&b)^(c^d)
        """
        return self.add_statement(other, '^')

    def __invert__(self):
        r"""
        Overloads the ~ operator to not a statement.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns a new statement that is the first statement logically
            not'ed.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: ~s
            ~(a&b)
        """
        exp = '~(' + self.__expression + ')'
        parse_tree, vars_order = logicparser.parse(exp)
        return BooleanFormula(exp, parse_tree, vars_order)

    def ifthen(self, other):
        r"""
        Returns two statements attached by the -> operator.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns a new statement that is the first statement logically
            ifthen'ed together.

       EXAMPLES:
           sage: import sage.logic.propcalc as propcalc
           sage: s = propcalc.formula("a&b")
           sage: f = propcalc.formula("c^d")
           sage: s.ifthen(f)
           (a&b)->(c^d)
        """
        return self.add_statement(other, '->')

    def iff(self, other):
        r"""
        Returns two statements attached by the <-> operator.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns a new statement that is the first statement logically
            ifandonlyif'ed together.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: f = propcalc.formula("c^d")
            sage: s.iff(f)
            (a&b)<->(c^d)
        """
        return self.add_statement(other, '<->')

    def __eq__(self, other):
        r"""
        Overloads the == operator to determine if two expressions
        are logically equivalent.

        INPUT:
            self -- the right hand side statement.
            other -- the left hand side statement.

        OUTPUT:
            Returns true if the left hand side is equivalent to the
            right hand side, and false otherwise.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: f = propcalc.formula("(a|b)&c")
            sage: g = propcalc.formula("c&(b|a)")
            sage: f == g
            True
            sage: g = propcalc.formula("a|b&c")
            sage: f == g
            False
        """
        return self.equivalent(other)

    def truthtable(self, start = 0, end = -1):
        r"""
        This function returns a truthtable object corresponding to the given
        statement.  Each row of the table corresponds to a binary number, with
        each variable associated to a column of the number, and taking on
        a true value if that column has a value of 1.  Please see the
        logictable module for details.  The function returns a table that
        start inclusive and end exclusive so truthtable(0, 2) will include
        row 0, but not row 2.

        INPUT:
            self -- the calling object.
            start -- an integer representing the row of the truth
                     table from which to start initialized to 0, which
                     is the first row when all the variables are
                     false.
            end -- an integer representing the last row of the truthtable
                   to be created.  It is initialized to the last row of the
                   full table.

        OUTPUT:
            Returns the truthtable (a 2-D array with the creating statement
            tacked on the front) corresponding to the statement.

        EXAMPLES:
        This example illustrates the creation of a statement.
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b|~(c|a)")
            sage: s.truthtable()
            a      b      c      value
            False  False  False  True
            False  False  True   False
            False  True   False  True
            False  True   True   False
            True   False  False  False
            True   False  True   False
            True   True   False  True
            True   True   True   True

        We can now create truthtable of rows 1 to 4, inclusive
            sage: s.truthtable(1, 5)
            a      b      c      value
            False  False  True   False
            False  True   False  True
            False  True   True   False
            True   False  False  False

        There should be no errors.

        NOTES:
            When sent with no start or end parameters, this is an
            exponential time function requiring O(2**n) time, where
            n is the number of variables in the expression.
        """
        max = 2 ** len(self.__vars_order)
        if(end < 0):
            end = max
        if(end > max):
            end = max
        if(start < 0):
            start = 0
        if(start > max):
            start = max
        keys, table = [], []
        vars = {}
        for var in self.__vars_order:
            vars[var] = False
            keys.insert(0, var)
        keys = list(keys)
        for i in range(start, end):
            j = 0
            row = []
            for key in keys:
                bit = self.get_bit(i, j)
                vars[key] = bit
                j += 1
                row.insert(0, bit)
            row.append(booleval.eval_formula(self.__tree, vars))
            table.append(row)
        keys.reverse()
        table = logictable.Truthtable(table, keys)
        return table

    def evaluate(self, var_values):
        r"""
        Evaluates a formula for the given input values.

        INPUT:
            self -- the calling object.
            var_values -- a dictionary containing pairs of
                              variables and their boolean values.
                              All variable must be present.

        OUTPUT:
            Return the evaluation of the formula with the given
            inputs, either True or False.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: f = propcalc.formula("a&b|c")
            sage: f.evaluate({'a':False, 'b':False, 'c':True})
            True
            sage: f.evaluate({'a':True, 'b':False, 'c':False})
            False
        """
        return booleval.eval_formula(self.__tree, var_values)

    def is_satisfiable(self):
        r"""
        Is_satisfiable determines if there is some assignment of
        variables for which the formula will be true.

        INPUT:
            self -- the calling object.

        OUTPUT:
            True if the formula can be satisfied, False otherwise.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: f = propcalc.formula("a|b")
            sage: f.is_satisfiable()
            True
            sage: g = f & (~f)
            sage: g.is_satisfiable()
            False
        """
        table = self.truthtable().get_table_list()
        for row in table[1:]:
            if row[-1] == True:
                return True
        return False

    def is_tautology(self):
        r"""
        Is_tautology determines if the formula is always
        true.

        INPUT:
            self -- the calling object.

        OUTPUT:
            True if the formula is a tautology, False otherwise.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: f = propcalc.formula("a|~a")
            sage: f.is_tautology()
            True
            sage: f = propcalc.formula("a&~a")
            sage: f.is_tautology()
            False
            sage: f = propcalc.formula("a&b")
            sage: f.is_tautology()
            False

        """
        return not (~self).is_satisfiable()

    def is_contradiction(self):
        r"""
        Is_contradiction determines if the formula is always
        false.

        INPUT:
            self -- the calling object.

        OUTPUT:
            True if the formula is a contradiction, False otherwise.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: f = propcalc.formula("a&~a")
            sage: f.is_contradiction()
            True
            sage: f = propcalc.formula("a|~a")
            sage: f.is_contradiction()
            False
            sage: f = propcalc.formula("a|b")
            sage: f.is_contradiction()
            False
        """
        return not self.is_satisfiable()

    def equivalent(self, other):
        r"""
        This function determines if two formulas are semantically
        equivalent.

        INPUT:
            self -- the calling object.
            other -- a boolformula instance.

        OUTPUT:
            True if the two formulas are logically equivalent, False
            otherwise.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: f = propcalc.formula("(a|b)&c")
            sage: g = propcalc.formula("c&(a|b)")
            sage: f.equivalent(g)
            True
            sage: g = propcalc.formula("a|b&c")
            sage: f.equivalent(g)
            False
        """
        return self.iff(other).is_tautology()

    def convert_cnf_table(self):
        r"""
        This function converts an instance of boolformula to conjunctive
        normal form. It does this by examining the truthtable of the formula,
        and thus takes `O(2^n)` time, where `n` is the number of variables.

        INPUT:

        - ``self`` -- the calling object.

        OUTPUT:

            An instance of :class:`BooleanFormula` with an identical truth
            table that is in conjunctive normal form.

        EXAMPLES::

            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a ^ b <-> c")
            sage: s.convert_cnf()
            sage: s
            (a|b|~c)&(a|~b|c)&(~a|b|c)&(~a|~b|~c)

        The methods :meth:`convert_cnf` and :meth:`convert_cnf_table` are aliases. ::

            sage: t = propcalc.formula("a ^ b <-> c")
            sage: t.convert_cnf_table(); t
            (a|b|~c)&(a|~b|c)&(~a|b|c)&(~a|~b|~c)
            sage: t == s
            True

        .. NOTE::

            This method creates the cnf parse tree by examining the logic
            table of the formula.  Creating the table requires `O(2^n)` time
            where `n` is the number of variables in the formula.
        """
        str = ''
        t = self.truthtable()
        table = t.get_table_list()
        vars = table[0]
        for row in table[1:]:
            if(row[-1] == False):
                str += '('
                for i in range(len(row) - 1):
                    if(row[i] == True):
                        str += '~'
                    str += vars[i]
                    str += '|'
                str = str[:-1] + ')&'
        self.__expression = str[:-1]
        #in case of tautology
        if(len(self.__expression) == 0):
            self.__expression = '(' + self.__vars_order[0] + '|~' + self.__vars_order[0] + ')'
        self.__tree, self.__vars_order = logicparser.parse(self.__expression)

    convert_cnf = convert_cnf_table

    def convert_cnf_recur(self):
        r"""
        This function converts an instance of boolformula to conjunctive
        normal form. It does this by applying a set of rules that are
        guaranteed to convert the formula.  Worst case the converted
        expression has an O(2^n) increase in size (and time as well), but if
        the formula is already in CNF (or close to) it is only O(n).

        INPUT:
            self -- the calling object.

        OUTPUT:
            An instance of boolformula with an identical truth table that is in
            conjunctive normal form.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a^b<->c")
            sage: s.convert_cnf_recur()
            sage: s
            (~a|a|c)&(~b|a|c)&(~a|b|c)&(~b|b|c)&(~c|a|b)&(~c|~a|~b)

        NOTES:
            This function can require an exponential blow up in space from the
            original expression.  This in turn can require large amounts of time.
            Unless a formula is already in (or close to) being in cnf convert_cnf()
            is typically preferred, but results can vary.
        """
        self.__tree = logicparser.apply_func(self.__tree, self.reduce_op)
        self.__tree = logicparser.apply_func(self.__tree, self.dist_not)
        self.__tree = logicparser.apply_func(self.__tree, self.dist_ors)
        self.convert_expression()

    def satformat(self):
        r"""
        This function returns the satformat representation of a boolean formula.
        See www.cs.ubc.ca/~hoos/SATLIB/Benchmarks/SAT/satformat.ps for a
        description of satformat.

        INPUT:
             self -- the calling object.

        OUTPUT:
            A string representing the satformat representation of this object.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: f = propcalc.formula("a&((b|c)^a->c)<->b")
            sage: f.convert_cnf()
            sage: f
            (a|~b|c)&(a|~b|~c)&(~a|b|~c)
            sage: f.satformat()
            'p cnf 3 0\n1 -2 3 0 1 -2 -3 \n0 -1 2 -3'

        NOTES:
            If the instance of boolean formula has not been converted to
            CNF form by a call to convert_cnf() or convert_cnf_recur()
            satformat() will call convert_cnf().  Please see the notes for
            convert_cnf() and convert_cnf_recur() for performance issues.
        """
        self.convert_cnf_table()
        s = ''
        vars_num = {}
        i = 0
        clauses = 0
        for e in self.__vars_order:
            vars_num[e] = str(i + 1)
            i += 1
        i = 0
        w = 1
        while i < len(self.__expression):
            c = self.__expression[i]
            if(c == ')'):
                clauses += 1
            if(c in '()|'):
                i += 1
                continue
            if(c == '~'):
                s += '-'
            elif(c == '&'):
                s += '0 '
            else:
                varname = ''
                while(i < self.__expression[i] not in '|) '):
                    varname += self.__expression[i]
                    i += 1
                s += vars_num[varname] + ' '
            if(len(s) >= (w * 15) and s[-1] != '-'):
                s += '\n'
                w += 1
            i += 1
        s = 'p cnf ' + str(len(self.__vars_order)) + ' ' + str(clauses) + '\n' + s
        return s[:-1]

#    def simplify(self):
#        r"""
#        This function uses the propcalc package to simplify an expression to
#        its minimal form.
#
#        INPUT:
#             self -- the calling object.
#
#        OUTPUT:
#            A simplified expression.
#
#        EXAMPLES:
#            sage: import sage.logic.propcalc as propcalc
#            sage: f = propcalc.formula("a&((b|c)^a->c)<->b")
#            sage: f.truthtable()
#            a      b      c      value
#            False  False  False  True
#            False  False  True   True
#            False  True   False  False
#            False  True   True   False
#            True   False  False  True
#            True   False  True   False
#            True   True   False  True
#            True   True   True   True
#            sage: f.simplify()
#            (~a&~b)|(a&~b&~c)|(a&b)
#            sage: f.truthtable()
#            a      b      c      value
#            False  False  False  True
#            False  False  True   True
#            False  True   False  False
#            False  True   True   False
#            True   False  False  True
#            True   False  True   False
#            True   True   False  True
#            True   True   True   True
#
#        NOTES:
#            If the instance of boolean formula has not been converted to
#            cnf form by a call to convert_cnf() or convert_cnf_recur()
#            satformat() will call convert_cnf().  Please see the notes for
#            convert_cnf() and convert_cnf_recur() for performance issues.
#        """
#        exp = ''
#        self.__tree = logicparser.apply_func(self.__tree, self.reduce_op)
#        plf = logicparser.apply_func(self.__tree, self.convert_opt)
#        wff = boolopt.PLFtoWFF()(plf) # convert to positive-normal form
#        wtd = boolopt.WFFtoDNF()
#        dnf = wtd(wff)
#        dnf = wtd.clean(dnf)
#        if(dnf == [] or dnf == [[]]):
#            exp = self.__vars_order[0] + '&~' + self.__vars_order[0] + ' '
#        opt = boolopt.optimize(dnf)
#        if(exp == '' and (opt == [] or opt == [[]])):
#            exp = self.__vars_order[0] + '|~' + self.__vars_order[0] + ' '
#        if(exp == ''):
#            for con in opt:
#                s = '('
#                for prop in con:
#                    if(prop[0] == 'notprop'):
#                       s += '~'
#                    s += prop[1] + '&'
#                exp += s[:-1] + ')|'
#        self.__expression = exp[:-1]
#        self.__tree, self.__vars_order = logicparser.parse(self.__expression)
#        return BooleanFormula(self.__expression, self.__tree, self.__vars_order)

    def convert_opt(self, tree):
        r"""
        This function can be applied to a parse tree to convert to the tuple
        form used by bool opt.  The expression must only contain '&', '|', and
        '~' operators.

        INPUT:
            self -- the calling object.
            tree -- a list of three elements corresponding to a branch of a
                    parse tree.
        OUTPUT:
            A tree branch that does not contain ^, ->, or <-> operators.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc, sage.logic.logicparser as logicparser
            sage: s = propcalc.formula("a&(b|~c)")
            sage: tree = ['&', 'a', ['|', 'b', ['~', 'c', None]]]
            sage: logicparser.apply_func(tree, s.convert_opt)
            ('and', ('prop', 'a'), ('or', ('prop', 'b'), ('not', ('prop', 'c'))))
        """
        if(type(tree[1]) is not TupleType and tree[1] != None):
            lval = ('prop', tree[1])
        else:
            lval = tree[1]
        if(type(tree[2]) is not TupleType and tree[2] != None):
            rval = ('prop', tree[2])
        else:
            rval = tree[2]
        if(tree[0] == '~'):
            return ('not', lval)
        if(tree[0] == '&'):
            op = 'and'
        if(tree[0] == '|'):
            op = 'or'
        return (op, lval, rval)

    def add_statement(self, other, op):
        r"""
        This function takes two statements and combines them
        together with the given operator.

        INPUT:
            self -- the left hand side statement object.
            other -- the right hand side statement object.
            op -- the character of the operation to be performed.

        OUTPUT:
            Returns a new statement that is the first statement attached to
            the second statement by the operator op.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: f = propcalc.formula("c^d")
            sage: s.add_statement(f, '|')
            (a&b)|(c^d)
            sage: s.add_statement(f, '->')
            (a&b)->(c^d)
        """
        exp = '(' + self.__expression + ')' + op + '(' + other.__expression + ')'
        parse_tree, vars_order = logicparser.parse(exp)
        return BooleanFormula(exp, parse_tree, vars_order)

    def get_bit(self, x, c):
        r"""
        This function returns bit c of the number x.  The 0 bit is the
        low order bit.  Errors should be handled gracefully by a return
        of false, and negative numbers x always return false while a
        negative c will index from the high order bit.

        INPUT:
            self -- the calling object.
            x -- an integer, the number from which to take the bit.
            c -- an integer, the bit number to be taken, where 0 is
                 the low order bit.

        OUTPUT:
            returns True if bit c of number x is 1, False otherwise.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a&b")
            sage: s.get_bit(2, 1)
            True
            sage: s.get_bit(8, 0)
            False

        It is not an error to have a bit out of range.
            sage: s.get_bit(64, 7)
            False

        Nor is it an error to use a negative number.
            sage: s.get_bit(-1, 3)
            False
            sage: s.get_bit(64, -1)
            True
            sage: s.get_bit(64, -2)
            False
        """
        bits = []
        while(x > 0):
            if(x % 2 == 0):
                b = False
            else:
                b = True
            x = int(x / 2)
            bits.append(b)
        if(c > len(bits) - 1):
            return False
        else:
            return bits[c]

    def reduce_op(self, tree):
        r"""
        This function can be applied to a parse tree to convert if-and-only-if,
        if-then, and xor operations to operations only involving and/or operations.

        INPUT:
            self -- the calling object.
            tree -- a list of three elements corresponding to a branch of a
                    parse tree.
        OUTPUT:
            A tree branch that does not contain ^, ->, or <-> operators.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc, sage.logic.logicparser as logicparser
            sage: s = propcalc.formula("a->b^c")
            sage: tree = ['->', 'a', ['^', 'b', 'c']]
            sage: logicparser.apply_func(tree, s.reduce_op)
            ['|', ['~', 'a', None], ['&', ['|', 'b', 'c'], ['~', ['&', 'b', 'c'], None]]]
        """
        if(tree[0] == '<->'):
            #parse tree for (~tree[1]|tree[2])&(~tree[2]|tree[1])
            new_tree = ['&', ['|', ['~', tree[1], None], tree[2]], \
                       ['|', ['~', tree[2], None], tree[1]]]
        elif(tree[0] == '^'):
            #parse tree for (tree[1]|tree[2])&~(tree[1]&tree[2])
            new_tree = ['&', ['|', tree[1], tree[2]], \
                       ['~', ['&', tree[1], tree[2]], None]]
        elif(tree[0] == '->'):
            #parse tree for ~tree[1]|tree[2]
            new_tree = ['|', ['~', tree[1], None], tree[2]]
        else:
            new_tree = tree
        return new_tree

    def dist_not(self, tree):
        r"""
        This function can be applied to a parse tree to distribute not operators
        over other operators.

        INPUT:
            self -- the calling object.
            tree -- a list of three elements corresponding to a branch of a
                    parse tree.
        OUTPUT:
            A tree branch that does not contain un-distributed nots.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc, sage.logic.logicparser as logicparser
            sage: s = propcalc.formula("~(a&b)")
            sage: tree = ['~', ['&', 'a', 'b'], None]
            sage: logicparser.apply_func(tree, s.dist_not) #long time
            ['|', ['~', 'a', None], ['~', 'b', None]]
        """
        if(tree[0] == '~' and type(tree[1]) is ListType):
            op = tree[1][0]
            if(op != '~'):
                if(op == '&'):
                    op = '|'
                else:
                    op = '&'
                new_tree = [op, ['~', tree[1][1], None], ['~', tree[1][2], None]]
                return logicparser.apply_func(new_tree, self.dist_not)
            else:
                #cancel double negative
                return tree[1][1]
        else:
            return tree

    def dist_ors(self, tree):
        r"""
        This function can be applied to a parse tree to distribute or over and.

        INPUT:
            self -- the calling object.
            tree -- a list of three elements corresponding to a branch of a
                    parse tree.
        OUTPUT:
            A tree branch that does not contain un-distributed ors.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc, sage.logic.logicparser as logicparser
            sage: s = propcalc.formula("(a&b)|(a&c)")
            sage: tree = ['|', ['&', 'a', 'b'], ['&', 'a', 'c']]
            sage: logicparser.apply_func(tree, s.dist_ors) #long time
            ['&', ['&', ['|', 'a', 'a'], ['|', 'b', 'a']], ['&', ['|', 'a', 'c'], ['|', 'b', 'c']]]

        """
        if(tree[0] == '|' and type(tree[2]) is ListType and tree[2][0] == '&'):
            new_tree = ['&', ['|', tree[1], tree[2][1]], ['|', tree[1], tree[2][2]]]
            return logicparser.apply_func(new_tree, self.dist_ors)
        if(tree[0] == '|' and type(tree[1]) is ListType and tree[1][0] == '&'):
           new_tree = ['&', ['|', tree[1][1], tree[2]], ['|', tree[1][2], tree[2]]]
           return logicparser.apply_func(new_tree, self.dist_ors)
        return tree

    def to_infix(self, tree):
        r"""
        This function can be applied to a parse tree to convert it from prefix to
        infix.

        INPUT:
            self -- the calling object.
            tree -- a list of three elements corresponding to a branch of a
                    parse tree.

        OUTPUT:
            A tree branch in infix form.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc, sage.logic.logicparser as logicparser
            sage: s = propcalc.formula("(a&b)|(a&c)")
            sage: tree = ['|', ['&', 'a', 'b'], ['&', 'a', 'c']]
            sage: logicparser.apply_func(tree, s.to_infix)
            [['a', '&', 'b'], '|', ['a', '&', 'c']]
        """
        if(tree[0] != '~'):
            return [tree[1], tree[0], tree[2]]
        return tree

    def convert_expression(self):
        r"""
        This function converts the string expression associated with an instance
        of boolformula to match with its tree representation after being converted
        to conjunctive normal form.

        INPUT:
            self -- the calling object.

        OUTPUT:
            None.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("a^b<->c")
            sage: s.convert_cnf_recur(); s  #long time
            (~a|a|c)&(~b|a|c)&(~a|b|c)&(~b|b|c)&(~c|a|b)&(~c|~a|~b)

        """
        ttree = self.__tree[:]
        ttree = logicparser.apply_func(ttree, self.to_infix)
        self.__expression = ''
        str_tree = str(ttree)
        open_flag = False
        put_flag = False
        i = 0
        for c in str_tree:
            if(i < len(str_tree) - 1):
                op = self.get_next_op(str_tree[i:])
            if(op == '|' and not open_flag):
                self.__expression += '('
                open_flag = True
            if(i < len(str_tree) - 2 and str_tree[i + 1] == '&' and open_flag):
                open_flag = False
                self.__expression += ')'
            if(str_tree[i:i + 4] == 'None'):
                i += 4
            if(i < len(str_tree) and str_tree[i] not in ' \',[]'):
                self.__expression += str_tree[i]
            i += 1
        if(open_flag == True):
            self.__expression += ')'

    def get_next_op(self, str):
        r"""
        This function returns the next operator in a string.

        INPUT:
            self -- the calling object.
            tree -- a string containing a logical expression.

        OUTPUT:
            The next operator in the string.

        EXAMPLES:
            sage: import sage.logic.propcalc as propcalc
            sage: s = propcalc.formula("f&p")
            sage: s.get_next_op("abra|cadabra")
            '|'
        """
        i = 0
        while(i < len(str) - 1 and str[i] != '&' and str[i] != '|'):
            i += 1
        return str[i]


def kbig(nums, a):
    if a==1:
        return max(nums)
    b = nums[0]
    c = max(nums)
    nums.remove(c)

    for i in range(0, len(nums)):

        for j in range(0, len(nums)):

            if abs(c - nums[j]) < b:
                b = abs(c - nums[j])
                d = nums[j]

        b=nums[0]

        c = d
        nums.remove(c)

        if i+2 == a:
            return c


import numpy as np
import operator
import sys
from scipy import spatial

def create_movie_user_dictionary(movie_data):
    movie_user_data = {}
    for i in range(0, 1682):
        for j in range(0, 943):
            user_list = []
            if not movie_data[j][i] == 0.0:
                user_list.append(j+1)
            movie_user_data[i+1] = user_list
    return movie_user_data
    # sys.stdout = open("movie_user_dict.txt", 'w')
    # print(movie_user_data)

if __name__ == '__main__':
    print("Program to predict the user rating for a movie he/she did not see")
    in_file = open("D:\OneDrive\IU\DataMining\Home works\HW1\ml-100k\/u.data", "r")
    next(in_file)
    np.set_printoptions(threshold=np.inf)
    mv_lens = np.zeros(943*1682).reshape(943, 1682)
    for line in in_file:
        line = list(line.split())
        user_id, item_id, rating = int(line[0]), int(line[1]), int(line[2])
        mv_lens[user_id-1][item_id-1] = rating

    movie_user_dict = create_movie_user_dictionary(mv_lens)
    for i in range(0, 943):
        for j in range(0, 1682):
            distance = {}
            if mv_lens[i][j] == 0.0:
                tot_rating = 0
                for user in movie_user_dict[j+1]:
                    distance[user] = (np.linalg.norm(mv_lens[user-1] - mv_lens[i]))  # Euclidean distance
                    distance[user] = (spatial.distance.cityblock(mv_lens[user-1] - mv_lens[i]))  # Manhattan distance
                sorted_distance = sorted(distance.items(), key=operator.itemgetter(1))
                if len(sorted_distance) > 10:
                    n = 20
                else:
                    n = len(sorted_distance)
                for user_distance_set in sorted_distance[:n]:
                    tot_rating += mv_lens[user_distance_set[0]-1][j]
                if n == 0:
                    predicted_rating = -1
                else:
                    predicted_rating = tot_rating/n
                mv_lens[i][j] = predicted_rating
    sys.stdout = open("final_output.txt", 'w')
    print(mv_lens)
    # in_file.close()


from collections import defaultdict
import torch
from torch_geometric.nn import HeteroConv
from torch_geometric.nn.conv import hetero_conv
from torch import nn
import math

# Taken from HeteroConv in torch_geometric.nn.conv
# added pos_dict to forward
class HeteroEGNN(HeteroConv):
    def __init__(self, convs, aggr='add', node_dim: int = 12):
        super().__init__(convs, aggr)
        self.node_dim = node_dim

    def forward(
        self,
        x_dict,
        pos_dict,
        edge_index_dict,
        *args_dict,
        **kwargs_dict,
    ):
        out_dict = defaultdict(list)
        for edge_type, edge_index in edge_index_dict.items():
            src, rel, dst = edge_type

            str_edge_type = '__'.join(edge_type)
            if str_edge_type not in self.convs:
                continue

            args = []
            for value_dict in args_dict:
                if edge_type in value_dict:
                    args.append(value_dict[edge_type])
                elif src == dst and src in value_dict:
                    args.append(value_dict[src])
                elif src in value_dict or dst in value_dict:
                    args.append(
                        (value_dict.get(src, None), value_dict.get(dst, None)))

            kwargs = {}
            for arg, value_dict in kwargs_dict.items():
                arg = arg[:-5]  # `{*}_dict`
                if edge_type in value_dict:
                    kwargs[arg] = value_dict[edge_type]
                elif src == dst and src in value_dict:
                    kwargs[arg] = value_dict[src]
                elif src in value_dict or dst in value_dict:
                    kwargs[arg] = (value_dict.get(src, None),
                                   value_dict.get(dst, None))

            conv = self.convs[str_edge_type]

            if self.node_dim == x_dict[src].shape[1]:
                out = conv(torch.hstack((pos_dict[src], x_dict[src])), edge_index, *args, **kwargs)
            else:
                out = conv(x_dict[src] , edge_index, *args, **kwargs)

            out_dict[dst].append(out)

        for key, value in out_dict.items():
            out_dict[key] = hetero_conv.group(value, self.aggr)

        return out_dict

class AttentionLayerNorm(nn.Module):
    def __init__(self, attention, embed_dim):
        super().__init__()
        self.attention = attention
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, *args):
        attn_output = self.attention(*args)
        return self.norm(attn_output)

    def pred_with_weights(self, *args):
        attn_output, attn_weights = self.attention.pred_with_weights(*args)
        return self.norm(attn_output), attn_weights

class SelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__init__()
        self.multihead_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True, dropout=dropout)

    def forward(self, x):
        attn_output, _ = self.multihead_attn(x, x, x)  # self-attention
        return attn_output

    def pred_with_weights(self, x):
        attn_output, attn_weights = self.multihead_attn(x, x, x)  # self-attention
        return attn_output, attn_weights

class CrossAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super().__init__()
        self.multihead_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True, dropout=dropout)

    def forward(self, query, kv):
        attn_output, _ = self.multihead_attn(query, kv, kv)  # cross-attention
        return attn_output

    def pred_with_weights(self, query, kv):
        attn_output, attn_weights = self.multihead_attn(query, kv, kv)  # cross-attention
        return attn_output, attn_weights

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        """
        Args
            d_model: Hidden dimensionality of the input.
            max_len: Maximum length of a sequence to expect.
        """
        super().__init__()

        # Create matrix of [SeqLen, HiddenDim] representing the positional encoding for max_len inputs
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)

        # register_buffer => Tensor which is not a parameter, but should be part of the modules state.
        # Used for tensors that need to be on the same device as the module.
        # persistent=False tells PyTorch to not add the buffer to the state dict (e.g. when we save the model)
        self.register_buffer("pe", pe, persistent=False)

    def forward(self, x):
        x = x + self.pe[:, : x.size(1)]
        return x

class ProteinEmbedding(nn.Module):
    def __init__(self,
                 residue_feature_dim,
                 protein_embedding_dim,
                 num_heads, num_layers,
                 dropout=0.1, max_prot_len=1000):
        super().__init__()
        self.protein_embedding_dim = protein_embedding_dim
        self.residue_scaler = nn.Sequential(
                nn.Linear(residue_feature_dim, protein_embedding_dim),
                nn.ReLU(),
                nn.Linear(protein_embedding_dim, protein_embedding_dim),
                nn.LayerNorm(protein_embedding_dim))
        self.positional_encoding = PositionalEncoding(self.protein_embedding_dim, max_len=max_prot_len)
        self.self_attn_layers = nn.ModuleList(
            [AttentionLayerNorm(
                SelfAttention(self.protein_embedding_dim, num_heads, dropout=dropout),
                self.protein_embedding_dim) for _ in range(num_layers)]
            )
        self.cross_attn_layers = nn.ModuleList(
            [AttentionLayerNorm(
                CrossAttention(self.protein_embedding_dim, num_heads, dropout=dropout),
                self.protein_embedding_dim) for _ in range(num_layers)]
            )

    def forward(self, residue_features):
        residue_features = self.positional_encoding(self.residue_scaler(residue_features))
        # Apply attention layers
        output_embedding = torch.ones((1, 1, self.protein_embedding_dim)).to(residue_features.device)
        for self_attn_layer, cross_attn_layer in zip(self.self_attn_layers, self.cross_attn_layers):
            residue_features = self_attn_layer(residue_features)
            output_embedding = cross_attn_layer(output_embedding, residue_features)
        return output_embedding.reshape((output_embedding.shape[0], self.protein_embedding_dim))


import logging
import unittest

import pytest
from pythoncommons.git_constants import ORIGIN

from tests.test_utilities import TestUtilities, Object, SANDBOX_REPO_DOWNSTREAM_HOTFIX
from yarndevtools.commands.backporter import Backporter
from yarndevtools.common.shared_command_utils import CommandType, RepoType
from yarndevtools.constants import TRUNK, BRANCH_3_1
from yarndevtools.yarn_dev_tools_config import DEFAULT_BASE_BRANCH

UPSTREAM_JIRA_ID = "YARN-123456: "
DOWNSTREAM_BRANCH = "cdh6x"
DOWNSTREAM_JIRA_ID = "CDH-1234"
UPSTREAM_REMOTE_NAME = "upstream"
FETCH = True

LOG = logging.getLogger(__name__)

# Commit should be in trunk, this is a prerequisite of the backporter
YARN_TEST_BRANCH = TRUNK
CHERRY_PICK_BASE_REF = TRUNK

class TestBackporter(unittest.TestCase):
    downstream_repo_wrapper = None
    downstream_utils = None
    upstream_utils = None
    upstream_repo = None
    log_dir = None
    sandbox_hadoop_repo_path = None

    @classmethod
    def setUpClass(cls):
        cls.upstream_utils = TestUtilities(cls, YARN_TEST_BRANCH)
        cls.upstream_utils.setUpClass(CommandType.BACKPORT, init_logging=True, console_debug=True)
        cls.upstream_utils.pull_to_trunk(ff_only=True)
        cls.upstream_repo = cls.upstream_utils.repo
        cls.upstream_repo_wrapper = cls.upstream_utils.repo_wrapper

        cls.downstream_utils = TestUtilities(cls, YARN_TEST_BRANCH)
        cls.downstream_utils.setUpClass(
            CommandType.BACKPORT, repo_postfix=SANDBOX_REPO_DOWNSTREAM_HOTFIX, init_logging=False
        )
        cls.downstream_utils.pull_to_trunk(ff_only=True)
        cls.downstream_repo = cls.downstream_utils.repo
        cls.downstream_repo_wrapper = cls.downstream_utils.repo_wrapper

        cls.full_ds_branch = f"{DOWNSTREAM_JIRA_ID}-{DOWNSTREAM_BRANCH}"
        cls.downstream_repo_wrapper.setup_committer_info("downstream_user", "downstream_email")
        # Setup debug logging of git commands
        cls.downstream_repo_wrapper.enable_debug_logging(full=True)

    @classmethod
    def tearDownClass(cls) -> None:
        TestUtilities.tearDownClass(cls.__name__, command_type=CommandType.BACKPORT)

    def setUp(self):
        self.upstream_utils.reset_and_checkout_existing_branch(YARN_TEST_BRANCH, pull=False)

        # THIS IS A MUST-HAVE!
        # Set up remote of upstream in the downstream repo
        self.downstream_repo_wrapper.add_remote(UPSTREAM_REMOTE_NAME, self.upstream_repo.git_dir)
        self.downstream_repo_wrapper.remove_branch(self.full_ds_branch, checkout_before_remove=TRUNK)

    @staticmethod
    def setup_args(cherry_pick_base_ref=None):
        args = Object()
        args.upstream_jira_id = UPSTREAM_JIRA_ID
        args.upstream_branch = DEFAULT_BASE_BRANCH
        args.downstream_jira_id = DOWNSTREAM_JIRA_ID
        args.downstream_branch = DOWNSTREAM_BRANCH
        args.no_fetch = not FETCH
        if cherry_pick_base_ref:
            args.downstream_base_ref = cherry_pick_base_ref
        return args

    def cleanup_and_checkout_branch_in_upstream_repo(self, branch=None, checkout_from=None, remove=True):
        self._cleanup_and_checkout_branch_internal(
            self.upstream_utils, RepoType.UPSTREAM, branch, checkout_from, remove
        )

    def cleanup_and_checkout_branch_in_downstream_repo(self, branch=None, checkout_from=None, remove=True):
        self._cleanup_and_checkout_branch_internal(
            self.downstream_utils, RepoType.DOWNSTREAM, branch, checkout_from, remove
        )

    def _cleanup_and_checkout_branch_internal(
        self, utils: TestUtilities, repo_type: RepoType, branch=None, checkout_from=None, remove=True
    ):
        if repo_type == RepoType.UPSTREAM:
            repo = self.upstream_repo
        elif repo_type == RepoType.DOWNSTREAM:
            repo = self.downstream_repo
        else:
            raise ValueError(
                "Unknown repo type '{}'. Only known repo types are: {}".format(repo_type, [e.value for e in RepoType])
            )

        if branch:
            utils.cleanup_and_checkout_test_branch(
                pull=False, branch=branch, checkout_from=checkout_from, remove=remove
            )
            self.assertEqual(branch, str(repo.head.ref))
        else:
            utils.cleanup_and_checkout_test_branch(pull=False, checkout_from=checkout_from)
            self.assertEqual(YARN_TEST_BRANCH, str(repo.head.ref))

    def test_with_uncommitted_should_raise_error(self):
        self.upstream_utils.add_some_file_changes(commit=False)
        args = self.setup_args()

        backporter = Backporter(args, self.upstream_repo_wrapper, self.downstream_repo_wrapper)
        self.assertRaises(ValueError, backporter.run)

    def test_with_committed_with_wrong_message_should_raise_error(self):
        self.cleanup_and_checkout_branch_in_upstream_repo()
        self.upstream_utils.add_some_file_changes(commit=True, commit_message_prefix="dummy")
        args = self.setup_args()

        backporter = Backporter(args, self.upstream_repo_wrapper, self.downstream_repo_wrapper)
        self.assertRaises(ValueError, backporter.run)

    def test_with_committed_with_good_message_remote_to_upstream_does_not_exist(self):
        self.cleanup_and_checkout_branch_in_upstream_repo()
        self.upstream_utils.add_some_file_changes(commit=True, commit_message_prefix=UPSTREAM_JIRA_ID)
        args = self.setup_args()

        # Intentionally remove remote
        self.downstream_repo_wrapper.remove_remote(UPSTREAM_REMOTE_NAME)

        backporter = Backporter(args, self.upstream_repo_wrapper, self.downstream_repo_wrapper)
        self.assertRaises(ValueError, backporter.run)

    def test_with_committed_with_good_message(self):
        self.cleanup_and_checkout_branch_in_upstream_repo()
        self.upstream_utils.add_some_file_changes(commit=True, commit_message_prefix=UPSTREAM_JIRA_ID)
        self.cleanup_and_checkout_branch_in_downstream_repo(branch=DOWNSTREAM_BRANCH)
        args = self.setup_args(cherry_pick_base_ref=CHERRY_PICK_BASE_REF)
        backporter = Backporter(args, self.upstream_repo_wrapper, self.downstream_repo_wrapper)
        backporter.run()

        expected_commit_msg = f"{DOWNSTREAM_JIRA_ID}: {UPSTREAM_JIRA_ID}test_commit"
        self.assertTrue(
            self.full_ds_branch in self.downstream_repo.heads,
            f"Created downstream branch does not exist: {self.full_ds_branch}",
        )
        self.downstream_utils.verify_commit_message_of_branch(
            self.full_ds_branch, expected_commit_msg, verify_cherry_picked_from=True
        )
        self.downstream_utils.verify_if_branch_is_moved_to_latest_commit(self.full_ds_branch)

    def test_backport_from_branch31(self):
        self.cleanup_and_checkout_branch_in_upstream_repo(branch=BRANCH_3_1, checkout_from=ORIGIN + "/" + BRANCH_3_1)
        self.upstream_utils.add_some_file_changes(commit=True, commit_message_prefix=UPSTREAM_JIRA_ID)
        args = self.setup_args(cherry_pick_base_ref=CHERRY_PICK_BASE_REF)
        args.upstream_branch = BRANCH_3_1

        backporter = Backporter(args, self.upstream_repo_wrapper, self.downstream_repo_wrapper)
        backporter.run()

        expected_commit_msg = f"{DOWNSTREAM_JIRA_ID}: {UPSTREAM_JIRA_ID}test_commit"
        self.assertTrue(
            self.full_ds_branch in self.downstream_repo.heads,
            f"Created downstream branch does not exist: {self.full_ds_branch}",
        )
        self.downstream_utils.verify_commit_message_of_branch(
            self.full_ds_branch, expected_commit_msg, verify_cherry_picked_from=True
        )
        self.downstream_utils.verify_if_branch_is_moved_to_latest_commit(self.full_ds_branch)

    def test_backport_relation_chain(self):
        self.cleanup_and_checkout_branch_in_upstream_repo(branch=BRANCH_3_1, checkout_from=ORIGIN + "/" + BRANCH_3_1)
        self.upstream_utils.add_some_file_changes(
            commit=True, commit_message_prefix=UPSTREAM_JIRA_ID, to_beginning=True
        )

        # Make some commits in downstream repo to simulate a relation chain
        relation_chain_branch_name = "my-relation-chain"
        self.cleanup_and_checkout_branch_in_downstream_repo(
            branch=relation_chain_branch_name, checkout_from=ORIGIN + "/" + TRUNK, remove=True
        )
        self.downstream_utils.add_some_file_changes(commit=True, commit_message_prefix="commit-1")
        self.downstream_utils.add_some_file_changes(commit=True, commit_message_prefix="commit-2")
        self.downstream_utils.add_some_file_changes(commit=True, commit_message_prefix="commit-3")

        args = self.setup_args(cherry_pick_base_ref=relation_chain_branch_name)
        args.upstream_branch = BRANCH_3_1

        backporter = Backporter(args, self.upstream_repo_wrapper, self.downstream_repo_wrapper)
        backporter.run()

        expected_commit_msg = f"{DOWNSTREAM_JIRA_ID}: {UPSTREAM_JIRA_ID}test_commit"
        self.assertTrue(
            self.full_ds_branch in self.downstream_repo.heads,
            f"Created downstream branch does not exist: {self.full_ds_branch}",
        )
        self.downstream_utils.verify_commit_message_of_branch(
            self.full_ds_branch, expected_commit_msg, verify_cherry_picked_from=True
        )
        self.downstream_utils.verify_if_branch_is_moved_to_latest_commit(self.full_ds_branch)
        self.downstream_utils.verify_if_branch_is_moved_to_latest_commit(relation_chain_branch_name)