from odoo import api, fields, models, _
from odoo.exceptions import UserError, ValidationError
from odoo.addons.delivery_ups.models.ups_request import UPSRequest, Package
from odoo.tools import pdf
import logging
_logger = logging.getLogger(__name__)

class ProviderUPS(models.Model):
    _inherit = 'delivery.carrier'

    def _get_ups_signature_required_type(self, order=None, picking=None, package=None):
        # '3' for Adult Sig.
        return '2'

    def _get_ups_signature_required(self, order=None, picking=None, package=None):
        if self.get_signature_required(order=order, picking=picking, package=package):
            return self._get_ups_signature_required_type(order=order, picking=picking, package=package)
        return False

    def _get_ups_is_third_party(self, order=None, picking=None):
        third_party_account = self.get_third_party_account(order=order, picking=picking)
        if third_party_account:
            if not third_party_account.delivery_type == 'ups':
                raise ValidationError('Non-UPS Shipping Account indicated during UPS shipment.')
            return True
        if order and order.ups_bill_my_account and order.ups_carrier_account:
            return True
        return False

    def _get_main_ups_account_number(self, order=None, picking=None):
        wh = None
        if order:
            wh = order.warehouse_id
        if picking:
            wh = picking.picking_type_id.warehouse_id
        if wh and wh.ups_shipper_number:
            return wh.ups_shipper_number
        return self.ups_shipper_number

    def _get_ups_account_number(self, order=None, picking=None):
        """
        Common hook to customize what UPS Account number to use.
        :return: UPS Account Number
        """
        # Provided by Hibou Odoo Suite `delivery_hibou`
        third_party_account = self.get_third_party_account(order=order, picking=picking)
        if third_party_account:
            if not third_party_account.delivery_type == 'ups':
                raise ValidationError('Non-UPS Shipping Account indicated during UPS shipment.')
            return third_party_account.name
        if order and order.ups_carrier_account:
            return order.ups_carrier_account
        if picking and picking.sale_id.ups_carrier_account:
            return picking.sale_id.ups_carrier_account
        if picking and picking.picking_type_id.warehouse_id.ups_shipper_number:
            return picking.picking_type_id.warehouse_id.ups_shipper_number
        return self.ups_shipper_number

    def _get_ups_carrier_account(self, picking):
        # 3rd party billing should return False if not used.
        account = self._get_ups_account_number(picking=picking)
        return account if account not in (self.ups_shipper_number, picking.picking_type_id.warehouse_id.ups_shipper_number) else False

    """
    Overrides to use Hibou Delivery methods to get shipper etc. and to add 'transit_days' to result.
    """
    def ups_rate_shipment(self, order):
        superself = self.sudo()
        # Hibou Shipping
        signature_required = superself._get_ups_signature_required(order=order)
        insurance_value = superself.get_insurance_value(order=order)
        currency = order.currency_id or order.company_id.currency_id
        insurance_currency_code = currency.name

        srm = UPSRequest(self.log_xml, superself.ups_username, superself.ups_passwd, superself._get_main_ups_account_number(order=order), superself.ups_access_number, self.prod_environment)
        ResCurrency = self.env['res.currency']
        max_weight = self.ups_default_packaging_id.max_weight
        packages = []
        total_qty = 0
        total_weight = 0
        for line in order.order_line.filtered(lambda line: not line.is_delivery and not line.display_type):
            total_qty += line.product_uom_qty
            total_weight += line.product_id.weight * line.product_qty

        if max_weight and total_weight > max_weight:
            total_package = int(total_weight / max_weight)
            last_package_weight = total_weight % max_weight

            for seq in range(total_package):
                packages.append(Package(self, max_weight,
                                        insurance_value=insurance_value, insurance_currency_code=insurance_currency_code, signature_required=signature_required))
            if last_package_weight:
                packages.append(Package(self, last_package_weight,
                                        insurance_value=insurance_value, insurance_currency_code=insurance_currency_code, signature_required=signature_required))
        else:
            packages.append(Package(self, total_weight,
                                    insurance_value=insurance_value, insurance_currency_code=insurance_currency_code, signature_required=signature_required))

        shipment_info = {
            'total_qty': total_qty  # required when service type = 'UPS Worldwide Express Freight'
        }

        if self.ups_cod:
            cod_info = {
                'currency': order.partner_id.country_id.currency_id.name,
                'monetary_value': order.amount_total,
                'funds_code': self.ups_cod_funds_code,
            }
        else:
            cod_info = None

        # Hibou Delivery
        shipper_company = self.get_shipper_company(order=order)
        shipper_warehouse = self.get_shipper_warehouse(order=order)
        recipient = self.get_recipient(order=order)
        date_planned = None
        if self.env.context.get('date_planned'):
            date_planned = self.env.context.get('date_planned')

        check_value = srm.check_required_value(shipper_company, shipper_warehouse, recipient, order=order)
        if check_value:
            return {'success': False,
                    'price': 0.0,
                    'error_message': check_value,
                    'warning_message': False}

        ups_service_type = order.ups_service_type or self.ups_default_service_type
        result = srm.get_shipping_price(
            shipment_info=shipment_info, packages=packages, shipper=shipper_company, ship_from=shipper_warehouse,
            ship_to=recipient, packaging_type=self.ups_default_packaging_id.shipper_package_code, service_type=ups_service_type,
            saturday_delivery=self.ups_saturday_delivery, cod_info=cod_info, date_planned=date_planned)

        if result.get('error_message'):
            return {'success': False,
                    'price': 0.0,
                    'error_message': _('Error:\n%s') % result['error_message'],
                    'warning_message': False}

        if order.currency_id.name == result['currency_code']:
            price = float(result['price'])
        else:
            quote_currency = ResCurrency.search([('name', '=', result['currency_code'])], limit=1)
            price = quote_currency._convert(
                float(result['price']), order.currency_id, order.company_id, order.date_order or fields.Date.today())

        # Hibou Delivery
        if self._get_ups_is_third_party(order=order):
            # Don't show delivery amount, if ups bill my account option is true
            price = 0.0

        return {'success': True,
                'price': price,
                'transit_days': result.get('transit_days', 0),
                'error_message': False,
                'warning_message': False}

    def ups_send_shipping(self, pickings):
        res = []
        superself = self.sudo()
        ResCurrency = self.env['res.currency']
        for picking in pickings:
            srm = UPSRequest(self.log_xml, superself.ups_username, superself.ups_passwd, superself._get_main_ups_account_number(picking=picking), superself.ups_access_number, self.prod_environment)
            # Hibou Delivery
            shipper_company = superself.get_shipper_company(picking=picking)
            shipper_warehouse = superself.get_shipper_warehouse(picking=picking)
            recipient = superself.get_recipient(picking=picking)
            currency = picking.sale_id.currency_id if picking.sale_id else picking.company_id.currency_id
            insurance_currency_code = currency.name

            picking_packages = picking.package_ids
            package_carriers = picking_packages.mapped('carrier_id')
            if package_carriers:
                # only ship ours
                picking_packages = picking_packages.filtered(lambda p: p.carrier_id == self and not p.carrier_tracking_ref)

            if not picking_packages:
                continue

            packages = []
            package_names = []
            if picking_packages:
                # Create all packages
                for package in picking_packages:
                    packages.append(Package(self, package.shipping_weight, quant_pack=package.packaging_id, name=package.name,
                                            insurance_value=superself.get_insurance_value(picking=picking, package=package), insurance_currency_code=insurance_currency_code, signature_required=superself._get_ups_signature_required(picking=picking, package=package)))
                    package_names.append(package.name)

            # what is the point of weight_bulk?
            # Create one package with the rest (the content that is not in a package)
            # if picking.weight_bulk:
            #     packages.append(Package(self, picking.weight_bulk))

            invoice_line_total = 0
            for move in picking.move_lines:
                invoice_line_total += picking.company_id.currency_id.round(move.product_id.lst_price * move.product_qty)

            shipment_info = {
                'description': superself.get_order_name(picking=picking),
                'total_qty': sum(sml.qty_done for sml in picking.move_line_ids),
                'ilt_monetary_value': '%d' % invoice_line_total,
                'itl_currency_code': self.env.user.company_id.currency_id.name,
                'phone': recipient.mobile or recipient.phone,
            }
            if picking.sale_id and picking.sale_id.carrier_id != picking.carrier_id:
                ups_service_type = picking.carrier_id.ups_default_service_type or picking.ups_service_type or superself.ups_default_service_type
            else:
                ups_service_type = picking.ups_service_type or superself.ups_default_service_type

            # Hibou Delivery
            ups_carrier_account = superself._get_ups_carrier_account(picking)

            if picking.carrier_id.ups_cod:
                cod_info = {
                    'currency': picking.partner_id.country_id.currency_id.name,
                    'monetary_value': picking.sale_id.amount_total,
                    'funds_code': superself.ups_cod_funds_code,
                }
            else:
                cod_info = None

            check_value = srm.check_required_value(shipper_company, shipper_warehouse, recipient, picking=picking)
            if check_value:
                raise UserError(check_value)

            package_type = picking_packages and picking_packages[0].packaging_id.shipper_package_code or self.ups_default_packaging_id.shipper_package_code
            srm.send_shipping(
                shipment_info=shipment_info, packages=packages, shipper=shipper_company, ship_from=shipper_warehouse,
                ship_to=recipient, packaging_type=package_type, service_type=ups_service_type, duty_payment=picking.carrier_id.ups_duty_payment,
                label_file_type=self.ups_label_file_type, ups_carrier_account=ups_carrier_account, saturday_delivery=picking.carrier_id.ups_saturday_delivery,
                cod_info=cod_info)
            result = srm.process_shipment()
            if result.get('error_message'):
                raise UserError(result['error_message'].__str__())

            order = picking.sale_id
            company = order.company_id or picking.company_id or self.env.company
            currency_order = picking.sale_id.currency_id
            if not currency_order:
                currency_order = picking.company_id.currency_id

            if currency_order.name == result['currency_code']:
                price = float(result['price'])
            else:
                quote_currency = ResCurrency.search([('name', '=', result['currency_code'])], limit=1)
                price = quote_currency._convert(
                    float(result['price']), currency_order, company, order.date_order or fields.Date.today())

            package_labels = []
            for track_number, label_binary_data in result.get('label_binary_data').items():
                package_labels = package_labels + [(track_number, label_binary_data)]

            carrier_tracking_ref = "+".join([pl[0] for pl in package_labels])
            logmessage = _("Shipment created into UPS
"
                           "Tracking Numbers: %s
"
                           "Packages: %s") % (carrier_tracking_ref, ','.join(package_names))
            if superself.ups_label_file_type != 'GIF':
                attachments = [('LabelUPS-%s.%s' % (pl[0], superself.ups_label_file_type), pl[1]) for pl in package_labels]
            if superself.ups_label_file_type == 'GIF':
                attachments = [('LabelUPS.pdf', pdf.merge_pdf([pl[1] for pl in package_labels]))]
            picking.message_post(body=logmessage, attachments=attachments)
            shipping_data = {
                'exact_price': price,
                'tracking_number': carrier_tracking_ref}
            res = res + [shipping_data]
            if self.return_label_on_delivery:
                self.ups_get_return_label(picking)
        return res

    def ups_rate_shipment_multi(self, order=None, picking=None, packages=None):
        if not packages:
            return self._ups_rate_shipment_multi_package(order=order, picking=picking)
        else:
            rates = []
            for package in packages:
                rates += self._ups_rate_shipment_multi_package(order=order, picking=picking, package=package)
            return rates

    def _ups_rate_shipment_multi_package(self, order=None, picking=None, package=None):
        superself = self.sudo()
        srm = UPSRequest(self.log_xml, superself.ups_username, superself.ups_passwd, superself._get_main_ups_account_number(order=order, picking=picking), superself.ups_access_number, self.prod_environment)
        ResCurrency = self.env['res.currency']
        max_weight = self.ups_default_packaging_id.max_weight
        packages = []
        if order:
            insurance_value = superself.get_insurance_value(order=order)
            signature_required = superself._get_ups_signature_required(order=order)
            currency = order.currency_id
            insurance_currency_code = currency.name
            company = order.company_id
            date_order = order.date_order or fields.Date.today()
            total_qty = 0
            total_weight = 0
            for line in order.order_line.filtered(lambda line: not line.is_delivery):
                total_qty += line.product_uom_qty
                total_weight += line.product_id.weight * line.product_qty

            if max_weight and total_weight > max_weight:
                total_package = int(total_weight / max_weight)
                last_package_weight = total_weight % max_weight

                for seq in range(total_package):
                    packages.append(Package(self, max_weight, insurance_value=insurance_value, insurance_currency_code=insurance_currency_code, signature_required=signature_required))
                if last_package_weight:
                    packages.append(Package(self, last_package_weight, insurance_value=insurance_value, insurance_currency_code=insurance_currency_code, signature_required=signature_required))
            else:
                packages.append(Package(self, total_weight, insurance_value=insurance_value, insurance_currency_code=insurance_currency_code, signature_required=signature_required))
        else:
            currency = picking.sale_id.currency_id if picking.sale_id else picking.company_id.currency_id
            insurance_currency_code = currency.name
            company = picking.company_id
            date_order = picking.sale_id.date_order or fields.Date.today() if picking.sale_id else fields.Date.today()
            # Is total quantity the number of packages or the number of items being shipped?
            if package:
                total_qty = 1
                packages = [Package(self, package.shipping_weight, insurance_value=superself.get_insurance_value(picking=picking, package=package), insurance_currency_code=insurance_currency_code, signature_required=superself._get_ups_signature_required(picking=picking, package=package))]
            elif picking.package_ids:
                # all packages....
                total_qty = len(picking.package_ids)
                packages = [Package(self, package.shipping_weight, insurance_value=superself.get_insurance_value(picking=picking, package=package), insurance_currency_code=insurance_currency_code, signature_required=superself._get_ups_signature_required(picking=picking, package=package)) for package in picking.package_ids.filtered(lambda p: not p.carrier_id)]
            else:
                total_qty = 1
                packages.append(Package(self, picking.shipping_weight or picking.weight, insurance_value=superself.get_insurance_value(picking=picking), insurance_currency_code=insurance_currency_code, signature_required=superself._get_ups_signature_required(picking=picking)))

        shipment_info = {
            'total_qty': total_qty  # required when service type = 'UPS Worldwide Express Freight'
        }

        if self.ups_cod:
            cod_info = {
                'currency': currency.name,
                'monetary_value': order.amount_total if order else picking.sale_id.amount_total,
                'funds_code': self.ups_cod_funds_code,
            }
        else:
            cod_info = None

        # Hibou Delivery
        shipper_company = self.get_shipper_company(order=order, picking=picking)
        shipper_warehouse = self.get_shipper_warehouse(order=order, picking=picking)
        recipient = self.get_recipient(order=order, picking=picking)
        date_planned = fields.Datetime.now()
        if self.env.context.get('date_planned'):
            date_planned = self.env.context.get('date_planned')

        check_value = srm.check_required_value(shipper_company, shipper_warehouse, recipient, order=order, picking=picking)
        if check_value:
            return [{'success': False,
                     'price': 0.0,
                     'error_message': check_value,
                     'warning_message': False,
                     }]
        # We now use Shop if we send multi=True
        ups_service_type = (order.ups_service_type or self.ups_default_service_type) if order else self.ups_default_service_type
        result = srm.get_shipping_price(
            shipment_info=shipment_info, packages=packages, shipper=shipper_company, ship_from=shipper_warehouse,
            ship_to=recipient, packaging_type=self.ups_default_packaging_id.shipper_package_code, service_type=ups_service_type,
            saturday_delivery=self.ups_saturday_delivery, cod_info=cod_info, date_planned=date_planned, multi=True)
        # Hibou Delivery
        is_third_party = self._get_ups_is_third_party(order=order, picking=picking)

        response = []
        for rate in result:
            if isinstance(rate, str):
                # assume error
                response.append({
                    'success': False, 'price': 0.0,
                    'error_message': _('Error:\n%s') % rate,
                    'warning_message': False,
                })
            elif rate.get('error_message'):
                _logger.error('UPS error: %s' % rate['error_message'])
                response.append({
                    'success': False, 'price': 0.0,
                    'error_message': _('Error:\n%s') % rate['error_message'],
                    'warning_message': False,
                })
            else:
                if currency.name == rate['currency_code']:
                    price = float(rate['price'])
                else:
                    quote_currency = ResCurrency.search([('name', '=', rate['currency_code'])], limit=1)
                    price = quote_currency._convert(
                        float(rate['price']), currency, company, date_order)

                if is_third_party:
                    # Don't show delivery amount, if ups bill my account option is true
                    price = 0.0

                service_code = rate['service_code']
                carrier = self.ups_find_delivery_carrier_for_service(service_code)
                if carrier:
                    response.append({
                        'carrier': carrier,
                        'package':  package or self.env['stock.quant.package'].browse(),
                        'success': True,
                        'price': price,
                        'error_message': False,
                        'warning_message': False,
                        'date_planned': date_planned,
                        'date_delivered': None,
                        'transit_days': rate.get('transit_days', 0),
                        'service_code': service_code,
                    })
        return response

    def ups_find_delivery_carrier_for_service(self, service_code):
        if self.ups_default_service_type == service_code:
            return self
        # arbitrary decision, lets find the same account number
        carrier = self.search([('ups_shipper_number', '=', self.ups_shipper_number),
                               ('ups_default_service_type', '=', service_code)
                               ], limit=1)
        return carrier


import csv
import json
import os
import re

import tact.processing.datetime_parser as datetime_parser
import tact.util.constants as constants
from tact.control.logging_controller import LoggingController as loggingController

logger = loggingController.get_logger(__name__)

def find_date_components(field_names):
    mn = dy = yr = "Not Found"

    date_search_pattern = re.compile("^date.*", re.IGNORECASE)
    month_search_pattern = re.compile("month|mnth(?=s| |$)", re.IGNORECASE)
    day_search_pattern = re.compile("day(?=s| |$)", re.IGNORECASE)
    year_search_pattern = re.compile("year(?=s| |$)", re.IGNORECASE)

    for current in field_names:
        date_match = date_search_pattern.match(current)
        if date_match:
            return {"date": date_match.group(0)}

        month_match = month_search_pattern.match(current)
        if month_match:
            mn = month_match.group(0)
            continue

        day_match = day_search_pattern.match(current)
        if day_match:
            dy = day_match.group(0)
            continue

        year_match = year_search_pattern.match(current)
        if year_match:
            yr = year_match.group(0)
            continue

    return {"year": yr, "month": mn, "day": dy}

def find_time_field(field_names):
    hr = min = sec = "Not Found"

    time_search_pattern = re.compile("^time.*", re.IGNORECASE)
    hour_search_pattern = re.compile("^h(?=our|r)", re.IGNORECASE)
    minute_search_pattern = re.compile("^min(?=ute|s|$)", re.IGNORECASE)
    second_search_pattern = re.compile("sec(?=ond|s|$)", re.IGNORECASE)

    for current in field_names:
        time_match = time_search_pattern.match(current)
        if time_match:
            return {"time": time_match.group(0)}

        hour_match = hour_search_pattern.match(current)
        if hour_match:
            hr = hour_match.string
            continue

        mintute_match = minute_search_pattern.match(current)
        if mintute_match:
            min = mintute_match.string
            continue

        second_match = second_search_pattern.match(current)
        if second_match:
            sec = second_match.string
            continue

    return {"hour": hr, "minute": min, "second": sec}

def process_file(input_path, input_encoding):

    with open(constants.PARSER_CONFIG_FILE_PATH) as json_file:
        config = json.load(json_file)

    # if isDirectory, get path to first file in dir, and proceed
    if config["isDirectory"]:
        # remove trailing slash
        if config["inputPath"].endswith("/"):
            config["inputPath"] = config["inputPath"].removesuffix("/")

        # gets list of files, ignoring hidden
        raw_files = [
            f for f in os.listdir(config["inputPath"]) if not f.startswith(".")
        ]
        # add the full path back in
        file_paths = list(
            map(lambda current: config["inputPath"] + "/" + current, raw_files)
        )

        input_path = file_paths[0]

        config["pathForPreview"] = input_path

        logger.info(
            "Input path is a directory, retrieving first file for preview: %s",
            input_path,
        )

    f = open(input_path, encoding=input_encoding)
    logger.debug("Opened file: %s", input_path)

    # # optionally skipping header rows
    # if settings_JSON['headerRow']:
    #     for i in range(settings_JSON['headerRow'] - 1):
    #         f.next()

    reader = csv.DictReader(f)
    # get field names and update
    field_names = reader.fieldnames
    config["fieldNames"] = field_names

    # get date & time fields
    config["dateFields"] = find_date_components(field_names)
    config["timeField"] = find_time_field(field_names)

    # THIS SHOULD NOT BE WRITING SETTINGS DIRECTLY :(
    # USE THE API or update_settings in controller
    # write out settings file
    with open(constants.CONFIG_FILE_PATHS["parser"], "w") as outfile:
        json.dump(config, outfile)
        logger.info("Settings written to: %s", constants.CONFIG_FILE_PATHS["parser"])

    return True

# open file
# grab a subset of rows
# grab one row for each unique date or time format
# get the length of all records
# call create_iso_time(csv_row, dateFields, timeField)
# output the results to a JSON file and return
def create_preview(config):
    logger.info("Generating preview")

    # parse preview file
    with open(config["pathForPreview"], encoding=config["inputFileEncoding"]) as f:
        reader = csv.DictReader(f)

        # skipping n rows to get to data
        next(reader)

        sample_JSON = {}
        sample_JSON["samples"] = []
        known_date_lengths = []
        known_time_lengths = []
        # TODO: add handling for multiple date fields
        # TODO: add option to start processing at arbitary line (in case data doesn't start on line #2)
        # take sum of characters in all date fields
        for csv_row in reader:
            # set lengths
            date_length = 0
            time_length = 0

            if config["dateFields"] != "Not Found":
                for field in config["dateFields"]:
                    if config["dateFields"][field] != "Not Found":
                        date_length += len(csv_row[config["dateFields"][field]])

            if config["timeField"] != "Not Found":
                for field in config["timeField"]:
                    if config["timeField"][field] != "Not Found":
                        time_length += len(csv_row[config["timeField"][field]])

            if (
                date_length not in known_date_lengths
                or time_length not in known_time_lengths
            ):
                current = {}
                # construct new JSON date object by looping through the fields in
                # config['dateFields']
                for key, value in config["dateFields"].items():
                    if value != "Not Found":
                        field_name = "Original_" + value
                        current[field_name] = csv_row[config["dateFields"].get(key)]

                for key, value in config["timeField"].items():
                    if value != "Not Found":
                        field_name = "Original_" + value
                        current[field_name] = csv_row[config["timeField"].get(key)]

                current["Transformation"] = datetime_parser.create_iso_time(
                    csv_row, config["dateFields"], config["timeField"]
                )
                sample_JSON["samples"].append(current)

                # add relevant field length to known values
                if date_length not in known_date_lengths:
                    known_date_lengths.append(date_length)

                if time_length not in known_time_lengths:
                    known_time_lengths.append(time_length)

    return sample_JSON


import math
import time
from copy import copy

import matplotlib.pyplot as plt
import random as rd
import string
import numpy as np

class Specimen:
    def __init__(self, sectors: dict, points: float):
        self.sectors = sectors
        self.points = points

    def __repr__(self):
        return f"Specimen({self.points}, {self.sectors})"

    def __str__(self):
        return f"Specimen({self.points}, {self.sectors})"

def gen_first_gen(n_sects, mach_names, n_first_gen) -> list:
    first_gen = []
    for i in range(n_first_gen):
        rd.shuffle(mach_names)
        letters = list(string.ascii_uppercase)[:n_sects]
        sectors = {letter: mach_names[it::n_sects] for it, letter in enumerate(letters)}
        first_gen.append(Specimen(sectors=sectors, points=0.0))
    return first_gen

def target(machines: dict, sectors: dict) -> tuple:
    """Average path"""
    avgs_of_paths = []
    for sector in sectors.keys():
        paths = []
        for i in range(len(sectors[sector])):
            cords0 = machines[sectors[sector][i-1]]
            cords1 = machines[sectors[sector][i]]
            path = math.sqrt((cords0[0] - cords1[0])**2 + (cords0[1] - cords1[1])**2)
            paths.append(path)
        avg_path = sum(paths)/len(paths)
        avgs_of_paths.append(avg_path)
    avg = sum(avgs_of_paths)/len(avgs_of_paths)
    std = float(np.std(avgs_of_paths))
    return avg, std

def fitness(avg: float, std: float) -> float:
    """The more points the better."""
    points = 1000.0/avg  # + std)
    # print(avg, std)
    return points

def selection(population: list, percentage_accepted: int) -> list:
    population.sort(key=lambda x: x.points, reverse=True)
    selected_population = population[:round(len(population) * percentage_accepted / 100)]
    # print([pts.points for pts in selected_population])
    return selected_population

def crossover(selected_population: list, new_pop_size: int, mach_names: list) -> list:
    small_sectors_len = min(len(x) for x in selected_population[0].sectors.values())
    big_sectors_len = max(len(x) for x in selected_population[0].sectors.values())
    small_sectors = []
    big_sectors = []

    for specimen in selected_population:
        for sector in specimen.sectors.values():
            if len(sector) == small_sectors_len:
                small_sectors.append(sector)
            elif len(sector) == big_sectors_len:
                big_sectors.append(sector)
            else:
                print("error")

    if not big_sectors:
        big_sectors = small_sectors

    n_sects = len(selected_population[0].sectors.keys())
    new_population = []
    for i in range(new_pop_size):
        # rd.shuffle(mach_names)
        letters = list(string.ascii_uppercase)[:n_sects]  # Remove "A" and "B"
        sectors = {}
        sector1 = rd.choice(small_sectors)
        sector2 = rd.choice(big_sectors)
        while [i for i in sector1 if i in sector2]:
            sector1 = rd.choice(small_sectors)
            sector2 = rd.choice(big_sectors)

        let1, let2 = "A", "A"
        while let1 == let2:
            let1 = rd.choice(letters)
            letters = [letter for letter in letters if letter is not let1]
            sectors[let1] = sector1
            let2 = rd.choice(letters)
            letters = [letter for letter in letters if letter is not let2]
            sectors[let2] = sector2
        used_machines = sector1 + sector2
        not_used_machines = [m_name for m_name in mach_names if m_name not in used_machines]
        other_sectors = {letter: not_used_machines[it::(n_sects-2)] for it, letter in enumerate(letters)}
        sectors.update(other_sectors)  # Mutation
        sector_keys = list(sectors.keys())
        sector_keys.sort()
        sectors = {i: sectors[i] for i in sector_keys}
        new_population.append(mutation(Specimen(sectors=sectors, points=0.0), mach_names=mach_names))
    return new_population

def mutation(specimen: Specimen, mach_names: list, chance=0.96) -> Specimen:
    if rd.random() >= chance * 0.9:
        rd.shuffle(mach_names)
        n_sects = len(specimen.sectors.keys())
        letters = list(string.ascii_uppercase)[:n_sects]
        sectors = {letter: mach_names[it::n_sects] for it, letter in enumerate(letters)}
        specimen = Specimen(sectors=sectors, points=0.0)
    elif rd.random() >= chance:
        rnd_sector1 = rd.choice(list(specimen.sectors.keys()))
        rnd_sector2 = rd.choice([key for key in specimen.sectors.keys() if key != rnd_sector1])
        rnd_machine1 = rd.choice(specimen.sectors[rnd_sector1])
        rnd_machine2 = rd.choice(specimen.sectors[rnd_sector2])
        specimen.sectors[rnd_sector1] = [machine for machine in specimen.sectors[rnd_sector1] if
                                         machine != rnd_machine1]
        specimen.sectors[rnd_sector2] = [machine for machine in specimen.sectors[rnd_sector2] if
                                         machine != rnd_machine2]
        specimen.sectors[rnd_sector1].append(rnd_machine2)
        specimen.sectors[rnd_sector2].append(rnd_machine1)
        rd.shuffle(specimen.sectors[rnd_sector1])
        rd.shuffle(specimen.sectors[rnd_sector2])
    return specimen

def select_best_specimen(population) -> Specimen:
    population.sort(key=lambda x: x.points, reverse=True)
    best_specimen = population[0]
    return best_specimen

def main() -> None:
    with open("data4.txt") as f:
        lines = f.readlines()
        machines_x = [int(number) for number in lines[1].split(",")]
        machines_y = [int(number) for number in lines[3].split(",")]
        machine_names = [name for name in lines[5].split(",")]
        n_sectors = int(lines[7])

    num_of_machines = len(machines_x)

    machine_names_scopy = machine_names.copy()
    machines_dict = {name: (machines_x[i], machines_y[i]) for i, name in enumerate(machine_names)}
    colormap = ['b', 'g', 'r', 'c', 'm', 'y', 'k', 'w']

    plt.ion()
    fig, ax = plt.subplots(figsize=(10, 8))
    factory_plot = ax.scatter(machines_x, machines_y, s=1000, marker='s')
    for j, txt in enumerate(machine_names):
        ax.annotate(txt, (machines_x[j], machines_y[j]), ha='center', va='center')
    ax.set_aspect('equal', adjustable='box')
    ax.set_xlim(left=min(machines_x)-10, right=max(machines_x)+10)
    ax.set_ylim(bottom=min(machines_y)-10, top=max(machines_y)+10)
    fig.canvas.draw()
    fig.canvas.flush_events()

    n_first_gen = 150
    n_next_population = 150
    percentage_accepted = 20

    population = gen_first_gen(n_sectors, machine_names, n_first_gen)

    best_ever = Specimen(sectors={}, points=0.0)
    best_ever_previous = copy(best_ever)
    start = time.time()
    counter = 0
    for i in range(10000):  # Main loop of algorithm
        for j, specimen in enumerate(population):
            avg, std = target(machines_dict, specimen.sectors)
            population[j].points = fitness(avg, std)

        selected_population = selection(population=population, percentage_accepted=percentage_accepted)

        best_spec = select_best_specimen(selected_population)
        if best_spec.points > best_ever.points:
            best_ever = best_spec
        if best_ever.points == best_ever_previous.points:
            counter += 1
        if counter >= num_of_machines*10:
            break
        best_ever_previous = copy(best_ever)
        print(f"{i}. Best specimen: {best_spec.points}")
        # print(f"{i}. Best ever:\n{best_ever.points}")
        print(f"Time: {time.time() - start} s")

        if i % 10 == 0:
            colors_dict = {}
            for j, sector in enumerate(best_ever.sectors.keys()):
                for m_name in best_ever.sectors[sector]:
                    colors_dict[m_name] = colormap[j]
            colors_list = [colors_dict[m_name] for m_name in machine_names_scopy]
            factory_plot.set_color(colors_list)
            fig.canvas.draw()
            fig.canvas.flush_events()

        new_population = crossover(selected_population, n_next_population, machine_names)
        population = new_population
        # time.sleep(0.5)

    selected_population = selection(population=population, percentage_accepted=percentage_accepted)
    best_spec = select_best_specimen(selected_population)
    if best_spec.points < best_ever.points:
        best_ever = best_spec

    print(f"Best ever: {best_ever.points}")

    input()

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-

plt.close("all")
clear_all()

import Volt_Imfunctions as im
from scipy.ndimage.filters import gaussian_filter
from scipy.ndimage.morphology import binary_fill_holes

from keras.models import load_model
from keras import backend as K
K.set_image_dim_ordering('th') # color channel first

pathname=r"D:\Takashi\SPIM\06152017";
fname="Fish3-3"

br_threshold1=120
cell_radius=5
rank_radius=int(cell_radius*1.5)
model = load_model(os.getcwd()+r'\\functions\\cell_model_32.hdf5')

imdir=pathname+"\\"+fname+"\\Registered\\"
ave=imread(imdir+"\\ave.tif")

image_radius=14 ## fixed value for deep learning datasets
mask=binary_fill_holes(ave>br_threshold1)
rank_image=gaussian_filter(im.imrank_dll(gaussian_filter(ave,cell_radius*0.2),rank_radius),cell_radius*0.2)
minima_image=im.local_minima_dll(rank_image,int(cell_radius*1.5))
overlay=np.tile(im.imNormalize(ave,95)[:,:,None],(1,1,3))/2
cell_centers=np.where((minima_image>0) & mask)
cell_centers_yx=np.array([cell_centers[0],cell_centers[1]]).T

overlay[cell_centers[0],cell_centers[1],0]=1
overlay[cell_centers[0],cell_centers[1],1]=0
overlay[cell_centers[0],cell_centers[1],2]=0

plt.figure(1)
plt.imshow(overlay)

v1=ave.mean()
v2=ave.std()
ave_pad=(np.random.randn(ave.shape[0]+image_radius*2,ave.shape[1]+image_radius*2)*np.sqrt(v2))+v1
ave_pad[image_radius:-image_radius,image_radius:-image_radius]=ave
rank_pad=np.zeros((ave.shape[0]+image_radius*2,ave.shape[1]+image_radius*2))
rank_pad[image_radius:-image_radius,image_radius:-image_radius]=rank_image

cell_r=cell_centers[0][np.where(cell_centers[0]>cell_radius*2)]
cell_r=cell_centers[0]
cell_c=cell_centers[1]

imlist=np.zeros((len(cell_r),1,image_radius*2+1,image_radius*2+1))
classlist=np.zeros((len(cell_r),))
for i in range(len(cell_r)):
    imlist[i,0,:,:]=im.imNormalize(ave_pad[cell_r[i]:(cell_r[i]+image_radius*2+1), \
          cell_c[i]:(cell_c[i]+image_radius*2+1)],99)[None,None,:,:]

scores=model.predict(imlist)
label=(scores[:,0]>0.5).astype('int')
label[np.where(label==0)]=2

inds1=np.where(label==1)[0]
inds2=np.where(label==2)[0]

disk=im.makeDisk(2)
disk_r=np.where(disk>0)[0]-2
disk_c=np.where(disk>0)[1]-2
cell_image=np.tile(im.imNormalize(ave,99.9)[:,:,None],(1, 1, 3))
cell_image_pad=np.zeros((cell_image.shape[0]+image_radius*2,cell_image.shape[1]+image_radius*2,3))
cell_image_pad[image_radius:-image_radius,image_radius:-image_radius,:]=cell_image
cell_image_pad2=cell_image_pad.copy()
for i in range(len(cell_r)):
    if label[i]==1:
        cell_image_pad[cell_r[i]+disk_r+image_radius,cell_c[i]+disk_c+image_radius,0]=1
        cell_image_pad[cell_r[i]+disk_r+image_radius,cell_c[i]+disk_c+image_radius,1]=0
        cell_image_pad[cell_r[i]+disk_r+image_radius,cell_c[i]+disk_c+image_radius,2]=0
    if label[i]==2:
        cell_image_pad[cell_r[i]+disk_r+image_radius,cell_c[i]+disk_c+image_radius,0]=0
        cell_image_pad[cell_r[i]+disk_r+image_radius,cell_c[i]+disk_c+image_radius,1]=1
        cell_image_pad[cell_r[i]+disk_r+image_radius,cell_c[i]+disk_c+image_radius,2]=1

plt.figure(2)
plt.imshow(cell_image_pad)

cell_inds1=im.proof_images(imlist,inds1)

cell_r=cell_r[cell_inds1]
cell_c=cell_c[cell_inds1]
cnum1=len(cell_inds1)

ROI_inds1=im.draw_periphery(ave,cell_r,cell_c,cell_radius*2)

for i in range(len(cell_inds1)):

    cell_image_pad2[ROI_inds1[i][0]+image_radius,ROI_inds1[i][1]+image_radius,0]=1
    cell_image_pad2[ROI_inds1[i][0]+image_radius,ROI_inds1[i][1]+image_radius,1]=0
    cell_image_pad2[ROI_inds1[i][0]+image_radius,ROI_inds1[i][1]+image_radius,2]=0

fig=plt.figure(3)
plt.subplot(121).imshow(im.imNormalize(ave_pad,99.9))
a2=plt.subplot(122)
xs=[]
ys=[]
while True:

    a2.cla()
    plt.subplot(122).imshow(cell_image_pad2)
    fig.canvas.draw()
    xy=fig.ginput(1)

    if not xy:
        break
    else:
        if (xy[0][0]<ave.shape[1]+image_radius) and (xy[0][1]<ave.shape[0]+image_radius) and (min(xy[0])>=image_radius):
            xs.append(int(xy[0][0])-image_radius)
            ys.append(int(xy[0][1])-image_radius)
            ROI=im.draw_periphery_single(ave,int(xy[0][1]-image_radius),int(xy[0][0]-image_radius),cell_radius*2)
            cell_image_pad2[ROI[0]+image_radius,ROI[1]+image_radius,0]=1
            cell_image_pad2[ROI[0]+image_radius,ROI[1]+image_radius,1]=0
            cell_image_pad2[ROI[0]+image_radius,ROI[1]+image_radius,2]=0

ROI_inds2=im.draw_periphery(ave,np.array(ys),np.array(xs),cell_radius*2)

for i in range(len(ys)):

    cell_image_pad2[ROI_inds2[i][0]+image_radius,ROI_inds2[i][1]+image_radius,0]=1
    cell_image_pad2[ROI_inds2[i][0]+image_radius,ROI_inds2[i][1]+image_radius,1]=0
    cell_image_pad2[ROI_inds2[i][0]+image_radius,ROI_inds2[i][1]+image_radius,2]=0

fig=plt.figure(4)
plt.subplot(121).imshow(im.imNormalize(ave_pad,99.9))
plt.subplot(122).imshow(cell_image_pad2)

ROI_list=ROI_inds1+ROI_inds2

np.save(imdir+"\\ROI_list.npy",ROI_list)

nROI=len(ROI_list)

plt.figure(5,figsize=(18,6))
plt.subplot(131).imshow(im.imNormalize(ave.astype('float'),99.9), cmap='gray')
img_color=np.tile(im.imNormalize(ave.astype('float'),99.9)[:,:,None],(1,1,3))
for n in range(nROI):
    inds=ROI_list[n]
    img_color[inds[0].astype('int'),inds[1].astype('int'),0]=1
    img_color[inds[0].astype('int'),inds[1].astype('int'),1]=0
    img_color[inds[0].astype('int'),inds[1].astype('int'),2]=0

plt.subplot(132).imshow(img_color)

ax=plt.subplot(133)
tmp=np.zeros(ave.shape)
for n in range(nROI):
    inds=ROI_list[n]
    tmp[inds[0].astype('int'),inds[1].astype('int')]=1

plt.imshow(tmp, cmap='gray')

for n in range(nROI):
    inds=ROI_list[n]
    location=(inds[0].mean(),inds[1].mean())
    ax.text(location[1],location[0],n,fontsize=12, color='r')

plt.savefig(imdir+"ROI.png")


from matplotlib import pyplot
import numpy
import scipy.stats
import seaborn
import pickle

import prob_spline
import test_common
import pandas as pd
import pickle
import joblib
from time import gmtime, strftime

def parallel_splines(to_be_run,file_name,N,Mos_Class=0,sigma=0,sample=0,combine_index=[],remove_index=[]):
	if Mos_Class==0:
		if N==1:
			output = to_be_run(file_name,sigma=sigma,sample=sample,combine_index=combine_index,remove_index=remove_index,seed=1)
		else:

			print('Start Time')
			print(strftime("%Y-%m-%d %H:%M:%S", gmtime()))
			with joblib.Parallel(n_jobs=-1) as parallel:
				output = parallel(joblib.delayed(to_be_run)(file_name,sigma=sigma,sample=sample,combine_index=combine_index,remove_index=remove_index,
					seed=j+1) for j in range(N))
			print('Finish Time')
			print(strftime("%Y-%m-%d %H:%M:%S", gmtime()))

	else:		#change function to take *kargs?
		if N==1:
			output = to_be_run(file_name,Mos_Class,sigma=sigma,sample=sample)
		else:
			print('Start Time')
			print(strftime("%Y-%m-%d %H:%M:%S", gmtime()))
			with joblib.Parallel(n_jobs=-1) as parallel:
				output = parallel(joblib.delayed(to_be_run)(file_name,Mos_Class,sigma=sigma,sample=sample) for j in range(N))

			print('Finish Time')
			print(strftime("%Y-%m-%d %H:%M:%S", gmtime()))
	return(output)

def get_splines(bc_file,bm_file,bc_sigma,bm_sigma,N,combine_index,remove_index):
	bc_splines = parallel_splines(prob_spline.HostSpline,bc_file,N,sigma=bc_sigma,sample=1,combine_index=combine_index,remove_index=remove_index)
	bm_splines = parallel_splines(prob_spline.BloodmealSpline,bm_file,N,sigma = bm_sigma,sample=1,combine_index=combine_index,remove_index=remove_index)
	if combine_index!=[]:
		with open('host_splines_sample_combine_index=%s.pkl' %str(combine_index), 'wb') as output:
			pickle.dump(bc_splines,output)
		with open('vectors_splines_sample_combine_index=%s.pkl' %str(combine_index), 'wb') as output:
			pickle.dump(bm_splines,output)
	elif remove_index!=[]:
		with open('host_splines_sample_remove_index=%s.pkl' %str(remove_index), 'wb') as output:
			pickle.dump(bc_splines,output)
		with open('vectors_splines_sample_remove_index=%s.pkl' %str(remove_index), 'wb') as output:
			pickle.dump(bm_splines,output)
	return()

def generate_splines_fun(combine_index=[],remove_index=[]):
	bc_file = "Days_BirdCounts.csv"
	msq_file = "Vector_Data(NoZeros).csv"
	bm_file = "Days_BloodMeal.csv"
	bc_sigma = 0		# for testing purposes
	#bc_sigma = pickle.load(open('Sigma_Vals_Cross_Validation_1mil_iter.pkl','rb'))
	#bm_sigma = 0.199
	bm_sigma= 0

	MosClass = prob_spline.MosConstant

	N = 2 # number of samples to be generated
	if not isinstance(bc_sigma,int):
		if combine_index != []:
			for j in combine_index:
				if j!=6:
					bc_sigma.pop(j)
		elif remove_index != []:
			bc_sigma.pop(remove_index)
	get_splines(bc_file,bm_file,bc_sigma,bm_sigma,N,combine_index,remove_index)
	#mos_curve = parallel_splines(prob_spline.MosCurve,msq_file,N,Mos_Class = prob_spline.MosConstant,sample=1)
	#with open('mos_curve_sample.pkl', 'wb') as output:
	#	pickle.dump(mos_curve,output)
	return()

if __name__ == '__main__':
	remove_index = [0,1,2,3,4,5,6]
	for x in remove_index:
		generate_splines_fun(remove_index=x)


from OpenGL.GL import *
from OpenGL.GL import shaders
from OpenGL import arrays
from OpenGL.arrays.arraydatatype import GLfloatArray
from OpenGL.GL.VERSION import GL_1_1

import glfw
import pygame
import glm
import math
import numpy
from ctypes import sizeof, c_float, c_void_p, c_uint, c_short
from collections import OrderedDict

from tools import CData

class ShaderData(object):
    SHA_PARAM = ('vao', 'program', 'texs')
    SHA_RENDER = SHA_TRANS = ('model', 'view', 'project', 'texture')
    SHA_CALLBACK = ('begin_render', 'end_render')
    SHA_NAMES = SHA_PARAM + SHA_RENDER + SHA_CALLBACK

    def __init__(self, **kwargs):
        self._id = None
        self._param = {}
        self._trans = {}
        self._render = {}
        self._callback = {}
        self.kwargs = kwargs

        self.store_trans(self.model_transform, self.view_transform, self.project_transform, self.texture_transform)

    def set_id(self, id):
        self._id = id

    def id(self):
        return self._id

    def _set(self, root, name, value):
        assert name in self.SHA_NAMES
        root[name] = value

    def _get(self, root, name):
        if isinstance(name, (tuple, list)):
            result = []
            for n in name:
                data = root.get(n)
                result.append(data)
            return result

        assert name in self.SHA_NAMES
        return root.get(name)

    def set_param(self, name, value):
        assert name not in self._param.keys()
        self._set(self._param, name, value)

    def get_param(self, name):
        return self._get(self._param, name)

    def store_param(self, *values):
        self._param.update(dict(zip(self.SHA_PARAM, values)))

    def set_trans(self, name, value):
        #assert name not in self._trans.keys()
        self._set(self._trans, name, value)

    def get_trans(self, name):
        return self._get(self._trans, name)

    def store_trans(self, *values):
        self._trans.update(dict(zip(self.SHA_TRANS, values)))

    def set_render(self, name, value):
        assert name not in self._render.keys()
        self._set(self._render, name, value)

    def get_render(self, name):
        return self._get(self._render, name)

    def store_render(self, *values):
        self._render.update(dict(zip(self.SHA_RENDER, values)))

    def set_callback(sell, name, value):
        assert name not in self._callback.keys()
        self._set(self._callback, name, value)

    def get_callback(self, name):
        return self._get(self._callback, name)

    def store_callback(self, *values):
        self._callback.update(dict(zip(self.SHA_CALLBACK, values)))

    def do_render(self, event):
        rstart = self.get_callback('begin_render')
        rend = self.get_callback('end_render')

        rstart(self)
        for n in self.SHA_RENDER:
            trans = self.get_trans(n)
            render = self.get_render(n)
            assert all((render, trans))

            mt = trans(event)
            render(self, mt)
        rend(self)

    def refresh(self, name, vbo):
        pass

    def model_transform(self, event, **kwargs):
        return glm.mat4(1)
        #self.store_id('model', model, shader_id)

    def view_transform(self, event, **kwargs):
        # eye指定相机位置，center指定相机指向目标位置，up指定viewUp向量

        # 默认从z轴正方向观看
        # eye = glm.vec3(0, 0, -1)
        # center = glm.vec3(0, 0, 0)
        # up = glm.vec3(0, 1, 0)
        #
        # view = glm.lookAt(eye, center, up)

        view = glm.mat4(1)

        return view

    def project_transform(self, event, **kwargs):
        #return glm.mat4(1)
        #self.store_id('project', proj, shader_id)

        # w, h = self.window_size
        # proj = glm.perspective(glm.radians(45), w / h, 1, 100)

        proj = glm.mat4(1)
        return proj

    def texture_transform(self, event, **kwargs):
        pass

class OpenglWrapper():

    def __init__(self, **kwargs):
        self.window_size = None
        self.window = None
        # self.program_id = None
        # self.vao_id = None
        # self.text_id = None

        self._shaders = []

    def create_window(self, w, h, event=None):
        if not glfw.init():
            return

        self.window_size = (w, h)

        glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
        glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
        glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
        glfw.window_hint(glfw.RESIZABLE, GL_FALSE)

        # Create a windowed mode window and its OpenGL context
        window = glfw.create_window(w, h, self.__class__.__name__, None, None)
        if not window:
            glfw.terminate()
            return

        # Make the window's context current
        glfw.make_context_current(window)
        glViewport(0, 0, w, h)

        if event:
            glfw.set_key_callback(window, event.key_callback)

            # 注册鼠标事件回调函数
            glfw.set_cursor_pos_callback(window, event.mouse_move_callback)
            # 注册鼠标滚轮事件回调函数
            glfw.set_scroll_callback(window, event.mouse_scroll_callback)
            # 鼠标捕获 停留在程序内
            glfw.set_input_mode(window, glfw.CURSOR, glfw.CURSOR_DISABLED)

        self.event = event
        self.window = window

    def create_shader(self, data, vformat, vcode, fcode, textures, cls_shader=ShaderData):
        vex_data = (data, vformat, None)
        datas = {'vertex': vex_data}
        return self.create_shader_instances(datas, vcode, fcode, textures, cls_shader)

    def create_shader_instances(self, datas, vcode, fcode, textures, cls_shader=ShaderData):
        shader = cls_shader()

        #parent = self
        #vertices = self.build_data(data)
        vao = self.build_vao(datas)
        program = self.build_shader(vcode, fcode)

        texs = []
        if textures:
            if isinstance(textures, (tuple, list)):
                for tex in textures:
                    t = self.build_texture(tex)
                    texs.append(t)
            else:
                t = self.build_texture(textures)
                texs.append(t)

        param = (vao, program, texs)
        if not all(param[:-1]):
            print(self.__class__.__name__, "create_shader failed,", "[param]", param)
            return

        shader.store_param(*param)
        #shader.store_trans(self.model_transform, self.view_transform, self.project_transform, self.texture_transform)
        shader.store_render(self.render_model, self.render_view, self.render_proj, self.render_texture)
        shader.store_callback(self._render_start, self._render_end)
        shader.set_id(len(self._shaders))

        self._shaders.append(shader)
        return shader

    def set_shader(self, type, name, value, shader_id=0):
        shader = self._shaders[shader_id]
        assert shader is not None

        if type == 'param':
            shader.set_param(name, value)
        elif type == 'trans':
            shader.set_trans(name, value)
        elif type == 'render':
            shader.set_render(name, value)
        else:
            raise Exception("Unknow shader type %s" % type)

    # def get_shader_param(self, name, shader_id=0):
    #     shader = self._shaders[shader_id]
    #     assert shader is not None
    #
    #     return shader.get_param(name)

    # def  build_data(self, data):
    #     return ArrayDatatype.asArray(data, GL_1_1.GL_FLOAT)

    def build_vbo(self, rdata, vformat, vdivisor=None, att_st=0):

        vertices = ArrayDatatype.asArray(rdata, GL_1_1.GL_FLOAT)

        vbo_id = glGenBuffers(1)

        glBindBuffer(GL_ARRAY_BUFFER, vbo_id)
        glBufferData(GL_ARRAY_BUFFER, len(vertices) * sizeof(c_float), vertices, GL_STATIC_DRAW)

        stride = sum(vformat) * sizeof(c_float)
        for i, size in enumerate(vformat):
            att_id = i + att_st
            offset = sum(vformat[:i]) * sizeof(c_float)
            glVertexAttribPointer(att_id, size, GL_FLOAT, GL_FALSE,
                                  stride, c_void_p(offset))
            glEnableVertexAttribArray(att_id)
            if vdivisor:
                glVertexAttribDivisor(att_id, vdivisor[i])

        atts = len(vformat)
        count = int(len(rdata) / sum(vformat))
        return dict(id=vbo_id, rdata=rdata, vformat=vformat, att_st=att_st, atts=atts, count=count)

    def update_vbo(self, vbo, buf_t=GL_DYNAMIC_DRAW):
        #GL_STATIC_DRAW
        #GL_STREAM_DRAW
        #GL_DYNAMIC_DRAW

        rdata = vbo['rdata']
        vbo_id = vbo['id']
        vertices = ArrayDatatype.asArray(rdata, GL_1_1.GL_FLOAT)

        glBindBuffer(GL_ARRAY_BUFFER, vbo_id)
        glBufferData(GL_ARRAY_BUFFER, len(vertices) * sizeof(c_float), vertices, buf_t)

    def build_vao(self, datas):
        vao_id = glGenVertexArrays(1)
        glBindVertexArray(vao_id)  # VAO record start

        vbos = {}
        att_st = 0
        for name, rdata in datas.items():
            vbo = self.build_vbo(*rdata, att_st)
            att_st += vbo['atts']
            vbos[name] = vbo

        glBindBuffer(GL_ARRAY_BUFFER, 0)
        glBindVertexArray(0)  # VAO record end

        return dict(id=vao_id, vbos=vbos)

    def build_shader(self, vcode, fcode):
        #vertex shader
        vertexShaderSource = vcode.encode()
        vertexShader = glCreateShader(GL_VERTEX_SHADER)
        glShaderSource(vertexShader, vertexShaderSource)
        glCompileShader(vertexShader)
        result = glGetShaderiv(vertexShader, GL_COMPILE_STATUS)
        if result == GL_FALSE:
            infoLog = glGetShaderInfoLog(vertexShader)
            print("vertex compile failed:", infoLog.decode('utf-8'))
            return
        else:
            print(vertexShaderSource, result)

        # fragment shader
        fragmentShaderSource = fcode.encode()
        fragmentShader = glCreateShader(GL_FRAGMENT_SHADER)
        glShaderSource(fragmentShader, fragmentShaderSource)
        glCompileShader(fragmentShader)
        result = glGetShaderiv(fragmentShader, GL_COMPILE_STATUS)
        if result == GL_FALSE:
            infoLog = glGetShaderInfoLog(fragmentShader)
            print("fragment compile failed::", infoLog.decode('utf-8'))
            return
        else:
            print(fragmentShaderSource, result)

        #link shader
        shaderProgram = glCreateProgram()
        glAttachShader(shaderProgram, vertexShader)
        glAttachShader(shaderProgram, fragmentShader)
        glLinkProgram(shaderProgram)
        result = glGetProgramiv(shaderProgram, GL_LINK_STATUS)
        if result == GL_FALSE:
            infoLog = glGetProgramInfoLog(shaderProgram)
            print("program link error:", infoLog.decode('utf-8'))
            return
        else:
            print("linked", result)

        glDetachShader(shaderProgram, vertexShader)
        glDetachShader(shaderProgram, fragmentShader)
        glDeleteShader(vertexShader)
        glDeleteShader(fragmentShader)

        return shaderProgram

    def build_texture(self, source):
        if not source:
            return None

        tex_id = glGenTextures(1)
        glBindTexture(GL_TEXTURE_2D, tex_id)
        # GL_REPEAT/GL_MIRRORED_REPEAT/GL_CLAMP_TO_EDGE/GL_CLAMP_TO_BORDER
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT)
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT)
        glTexParameteri(GL_TEXTURE_2D,
                        GL_TEXTURE_MAG_FILTER, GL_LINEAR)
        glTexParameteri(GL_TEXTURE_2D,
                        GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        # glTexParameteri(GL_TEXTURE_2D,
        #                 GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR)

        # load texture image
        pygame.init()
        # pygame.display.set_mode((w, h), pygame.DOUBLEBUF | pygame.OPENGL)
        img = pygame.image.load(source)
        if not img:
            return

        bitsize = img.get_bitsize()
        if bitsize == 32:
            src_format = 'RGBA'
            opg_format = GL_RGBA
        elif bitsize == 24:
            src_format = 'RGB'
            opg_format = GL_RGB
        else:
            print(self.__class__.__name__, "texture source not support", source, bitsize)
            return

        data = pygame.image.tostring(img, src_format, 1)
        img_width = img.get_width()
        img_height = img.get_height()
        out_format = GL_RGB
        glTexImage2D(GL_TEXTURE_2D, 0, out_format, img_width, img_height, 0, opg_format, GL_UNSIGNED_BYTE, data)
        glGenerateMipmap(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, 0)

        pygame.quit()

        return tex_id

    # def model_transform(self, shader):
    #     return glm.mat4(1)
    #     #self.store_id('model', model, shader_id)
    #
    # def view_transform(self, shader):
    #     return glm.mat4(1)
    #     #self.store_id('view', view, shader_id)
    #
    # def project_transform(self, shader):
    #     return glm.mat4(1)
    #     #self.store_id('project', proj, shader_id)
    #
    # def texture_transform(self, shader):
    #     tex = shader.get_param('texture')
    #     if tex is None:
    #         return
    #
    #     return tex

    def render_model(self, shader, model):
        # result = self.get_id(['program', 'model'], shader_id)
        # if not all(result):
        #     return
        #
        # program, model = result
        program = shader.get_param('program')
        glUniformMatrix4fv(glGetUniformLocation(program, "model"),
                           1, GL_FALSE, glm.value_ptr(model))

    def render_view(self, shader, view):
        # result = self.get_id(['program', 'view'], shader_id)
        # if not all(result):
        #     return
        #
        # program, view = result
        program = shader.get_param('program')
        glUniformMatrix4fv(glGetUniformLocation(program, 'view'),
                    1, GL_FALSE, glm.value_ptr(view))

    def render_proj(self, shader, project):
        # result = self.get_id(['program', 'project'], shader_id)
        # if not all(result):
        #     return
        #
        # program, project = result
        program = shader.get_param('program')
        glUniformMatrix4fv(glGetUniformLocation(program, 'projection'),
                           1, GL_FALSE, glm.value_ptr(project))

    def render_texture(self, shader, texture):
        # result = self.get_id(['program', 'texture'])
        # if not all(result):
        #     print(self.__class__.__name__, "render_texture vars empty", result)
        #     return
        #
        # program, tex = result

        program, texs= shader.get_param(['program', 'texs'])
        for i, tex in enumerate(texs):
            if tex:
                suffix = str(i)
                tag = globals()['GL_TEXTURE'+ suffix]   #GL_TEXTURE0 ...
                glActiveTexture(tag)
                glBindTexture(GL_TEXTURE_2D, tex)
                if i == 0:
                    suffix_f = ''
                else:
                    suffix_f = suffix
                t = glGetUniformLocation(program, "tex" + suffix_f) #tex, tex1, tex2 in program
                glUniform1i(t, i)  # 设置纹理单元为0号

    # def _render_start(self, shader):
    #     result = shader.get_param(['program', 'vao'])
    #     if not all(result):
    #         print(self.__class__.__name__, "_render_start var empty", result)
    #         return
    #
    #     program, vao_id = result
    #
    #     glBindVertexArray(vao_id)
    #     glUseProgram(program)

    # def _render_end(self, shader):
    #     result = self.get_shader_param(['vertices', 'vformat'])
    #     if not all(result):
    #         print(self.__class__.__name__, "_render_end var empty", result)
    #         return
    #
    #     vertices, vformat = result
    #
    #     # 使用索引绘制
    #     # glDrawElements(GL_TRIANGLES, len(indices), GL_UNSIGNED_SHORT, c_void_p(0))
    #     glDrawArrays(GL_TRIANGLES, 0, int(len(vertices) / sum(vformat)))

    def _render_start(self, shader):
        program, vao = shader.get_param(['program', 'vao'])
        vbos = vao['vbos']

        for name, vbo in vbos.items():
            result = shader.refresh(name, vbo)
            if result:
                self.update_vbo(vbo, GL_DYNAMIC_DRAW)

        glBindVertexArray(vao['id'])
        glUseProgram(program)

    def _render_end(self, shader):
        # 使用索引绘制
        vao = shader.get_param('vao')

        vbos = vao['vbos']
        verx = vbos.get('vertex')
        inst = vbos.get('inst')
        if not inst:
            dup_count = 0
        else:
            dup_count =inst['count']

        count = verx['count']

        if dup_count:
            glDrawArraysInstanced(GL_TRIANGLES, 0, count, dup_count)
        else:
            glDrawArrays(GL_TRIANGLES, 0, count)

    def render(self, shader):
        shader.do_render(self.event)

    def run(self):
        # Enable depth test
        glEnable(GL_DEPTH_TEST)
        # Accept fragment if it closer to the camera than the former one
        #glDepthFunc(GL_LESS)
        # Cull triangles which normal is not towards the camera
        #glEnable(GL_CULL_FACE)

        lastFrame = 0
        while not glfw.window_should_close(self.window):
            # Render here, e.g. using pyOpenGL
            glClearColor(0, 0, 0, 1.0)
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

            for shader in self._shaders:
                self.render(shader)

            # Swap front and back buffers
            glfw.swap_buffers(self.window)

            # Poll for and process events
            #glfw.poll_events()
            if self.event:
                currentFrame = glfw.get_time()
                deltaTime = currentFrame - lastFrame
                lastFrame = currentFrame

                # Poll for and process events
                glfw.poll_events()
                self.event.do_movement(deltaTime)
            else:
                glfw.poll_events()

            glBindVertexArray(0)
            glUseProgram(0)

        glfw.terminate()


# Original way to create a basic screen
# import pygame
# pygame.init()
#
# win = pygame.display.set_mode((500, 500))
#
# pygame.display.set_caption("First Game")
#
# x = 50
# y = 50
# width = 40
# height = 60
# vel = 5
#
# run = True
# while run:
#     pygame.time.delay(100)
#
#     for event in pygame.event.get():
#         if event.type == pygame.QUIT:
#             run = False
#
# pygame.quit()

# The link below shows and explains how to drop spheres with gravity playing a role with Pymunk and Pygames interacting.
# https://github.com/viblo/pymunk/blob/08fb141b81c0240513fc16e276d5ade5b0506512/docs/html/_sources/tutorials/SlideAndPinJoint.rst.txt

import sys, random
import pygame
from pygame.locals import *
import pymunk
import pymunk.pygame_util

import math, sys, random

import pygame
import re
from pygame.locals import *
from pygame.color import *
from pygame import mouse
import pymunk
from pymunk import Vec2d
import pymunk.pygame_util

object_draging = False  # object refers to any of the circles being dropped. This boolean indicates whether object is being dragged (mouse down) or not (mouse up)
square_draging = False  # refers to square object being dragged
ramp_dragging = False  # refers to ramp object being dragged
flipped_ramp_dragging = False
runonce = False

#testing
def draw_collision(arbiter, space, data):
    for c in arbiter.contact_point_set.points:
        r = max(3, abs(c.distance * 5))
        r = int(r)

        p = pymunk.pygame_util.to_pygame(c.point_a, data["surface"])
        pygame.draw.circle(data["surface"], THECOLORS["black"], p, r, 1)

def main():
    global contact
    global shape_to_remove
    global object_draging
    global square_draging
    global ramp_dragging
    global flipped_ramp_dragging
    global screen
    global numGrav
    numGrav = 1
    global balls
    global squares
    global ramps
    global flipped_ramps
    global space

    pygame.init()
    screen = pygame.display.set_mode((600, 600))
    clock = pygame.time.Clock()
    running = True

    ### Physics stuff
    global space
    space = pymunk.Space()  # creates a space for the physics

    draw_options = pymunk.pygame_util.DrawOptions(screen)
    # disable the build in debug draw of collision point since we use our own code.
    draw_options.flags = draw_options.flags ^ pymunk.pygame_util.DrawOptions.DRAW_COLLISION_POINTS
    ## Balls
    balls = []
    squares = []
    ramps = []
    flipped_ramps = []

    ### walls
    static_lines = [pymunk.Segment(space.static_body, (0, 5), (600, 5), 5)
        , pymunk.Segment(space.static_body, (595, 600), (595, 0), 5), pymunk.Segment(space.static_body, (0, 596), (600, 596), 5),
                    pymunk.Segment(space.static_body, (4, 0), (4, 600), 5)
                    ]
    for l in static_lines:
        l.friction = 0.5
    space.add(static_lines)

    # ticks_to_next_ball = 10

    ch = space.add_collision_handler(0, 0)
    ch.data["surface"] = screen
    ch.post_solve = draw_collision
    w, h = pygame.display.get_surface().get_size()

    while running:
        global ball
        global square
        global index_dragging
        global square_index_dragging
        global ramp_index_dragging
        global flipped_ramp_index_dragging
        for event in pygame.event.get():
            if event.type == QUIT:
                running = False
            elif event.type == KEYDOWN and event.key == K_ESCAPE:
                running = False
            # Added interactive stuff starting here:
            elif event.type == MOUSEBUTTONDOWN:
                if event.button == 1:  # When event.button == 1, that's referring to a left click

                    # 1 - left click
                    #
                    # 2 - middle click
                    #
                    # 3 - right click
                    #
                    # 4 - scroll up
                    #
                    # 5 - scroll down
                    index = -1
                    squareIndex = -1
                    ramp_index = -1
                    flipped_ramp_index = -1
                    for ball in balls:
                        surface = pygame.Surface(screen.get_size())
                        mouse_x, mouse_y = event.pos
                        index += 1
                        if (math.sqrt(math.pow(ball.body.position.x - mouse_x, 2) + math.pow((600-ball.body.position.y) - mouse_y, 2))) <= 25:
                            index_dragging = index  # saves the index of the ball that is clicked
                            space.gravity = (0, 0)
                            object_draging = True

                    for square in squares:
                        mouse_x, mouse_y = event.pos
                        squareIndex += 1
                        if mouse_x + 50 >= square.body.position.x and mouse_x <= square.body.position.x + 50 and mouse_y + 50 >= (600-square.body.position.y) and mouse_y <= (600-square.body.position.y) + 50:
                            square_index_dragging = squareIndex  # saves the index of the ball that is clicked
                            space.gravity = (0, 0)
                            square_draging = True
                            #mouse_x, mouse_y = event.pos

                    for ramp in ramps:
                        mouse_x, mouse_y = event.pos
                        ramp_index += 1
                        if mouse_x >= ramp.body.position.x and mouse_x <= ramp.body.position.x + 50 and mouse_y <= (600 - ramp.body.position.y) and mouse_y + 50 >= (600-ramp.body.position.y):
                            ramp_index_dragging = ramp_index
                            space.gravity = (0, 0)
                            ramp_dragging = True

                    for ramp in flipped_ramps:
                        mouse_x, mouse_y = event.pos
                        flipped_ramp_index += 1
                        print(ramp.body.position.x)
                        print((600-ramp.body.position.y))
                        if mouse_x <= ramp.body.position.x and mouse_x + 50 >= ramp.body.position.x and mouse_y <= (600 - ramp.body.position.y) and mouse_y + 50 >= (600-ramp.body.position.y):
                            flipped_ramp_index_dragging = flipped_ramp_index
                            space.gravity = (0, 0)
                            flipped_ramp_dragging = True

                    mouse_pos = mouse.get_pos()
                    x = pygame.mouse.get_pos()[0]
                    y = pygame.mouse.get_pos()[1]
                    if rect.collidepoint(mouse_pos):
                        if space.gravity == (0.0, 0):
                            space.gravity = (0.0, -980.0)  # sets the gravity of the space
                        else:
                            space.gravity = (0.0, 0)  # sets the gravity of the space
                    else:
                        if runonce is False and (object_draging is False) and (square_draging is False) and (ramp_dragging is False) and (flipped_ramp_dragging is False):
                            pos = pygame.mouse.get_pos()
                            final = str(pos)
                            x = final[final.find('(') + len('('):final.rfind(',')]
                            y = final[final.find(',') + len(','):final.rfind(')')]
                            add_circle(0.1, 25, x, y, 0.5, len(balls))

            elif event.type == pygame.MOUSEBUTTONUP:
                if event.button == 1:
                    object_draging = False
                    square_draging = False
                    ramp_dragging = False
                    flipped_ramp_dragging = False

            elif event.type == pygame.MOUSEMOTION:
                if object_draging:
                    mouse_x, mouse_y = event.pos
                    offset_x = balls[index_dragging].body.position.x - mouse_x
                    offset_y = balls[index_dragging].body.position.y - mouse_y
                    balls[index_dragging].body.position.x = mouse_x + offset_x
                    balls[index_dragging].body.position.y = mouse_y + offset_y
                    space.remove(balls[index_dragging])
                    balls.remove(balls[index_dragging])
                    add_circle(0.1, 25, mouse_x, mouse_y, 0.5, index_dragging)

                if square_draging:
                    mouse_x, mouse_y = event.pos
                    offset_x = squares[square_index_dragging].body.position.x - mouse_x
                    offset_y = squares[square_index_dragging].body.position.y - mouse_y
                    squares[square_index_dragging].body.position.x = mouse_x + offset_x
                    squares[square_index_dragging].body.position.y = mouse_y + offset_y
                    space.remove(squares[square_index_dragging])
                    squares.remove(squares[square_index_dragging])
                    add_square(0.1, 50.0, 50.0, mouse_x, mouse_y, 0.5, square_index_dragging)

                if ramp_dragging:
                    mouse_x, mouse_y = event.pos
                    offset_x = ramps[ramp_index_dragging].body.position.x - mouse_x
                    offset_y = ramps[ramp_index_dragging].body.position.y - mouse_y
                    ramps[ramp_index_dragging].body.position.x = mouse_x + offset_x
                    ramps[ramp_index_dragging].body.position.y = mouse_y + offset_y
                    space.remove(ramps[ramp_index_dragging])
                    ramps.remove(ramps[ramp_index_dragging])
                    add_ramp(100, 45, mouse_x, mouse_y, 0.1, ramp_index_dragging)

                if flipped_ramp_dragging:
                    mouse_x, mouse_y = event.pos
                    offset_x = flipped_ramps[flipped_ramp_index_dragging].body.position.x - mouse_x
                    offset_y = flipped_ramps[flipped_ramp_index_dragging].body.position.y - mouse_y
                    flipped_ramps[flipped_ramp_index_dragging].body.position.x = mouse_x + offset_x
                    flipped_ramps[flipped_ramp_index_dragging].body.position.y = mouse_y + offset_y
                    space.remove(flipped_ramps[flipped_ramp_index_dragging])
                    flipped_ramps.remove(flipped_ramps[flipped_ramp_index_dragging])
                    add_ramp_flipped(100, 45, mouse_x, mouse_y, 0.1, flipped_ramp_index_dragging)

            # Interactive stuff ends here

            elif event.type == KEYDOWN and event.key == K_p:
                pygame.image.save(screen, "contact_with_friction.png")
            elif event.type == KEYDOWN and event.key == K_s:
                pos = pygame.mouse.get_pos()
                final = str(pos)
                x = final[final.find('(') + len('('):final.rfind(',')]
                y = final[final.find(',') + len(','):final.rfind(')')]
                add_ramp(100, 45, x, y, 0.1, len(ramps))
            elif event.type == KEYDOWN and event.key == K_d:
                pos = pygame.mouse.get_pos()
                final = str(pos)
                x = final[final.find('(') + len('('):final.rfind(',')]
                y = final[final.find(',') + len(','):final.rfind(')')]
                add_ramp_flipped(100, 45, x, y, 0.1, len(ramps))

            # global runonce

            # adds a square when the 'A' key is pressed; change around later for better UI
            if event.type == KEYDOWN and event.key == K_a:
                pos = pygame.mouse.get_pos()
                final = str(pos)
                x = final[final.find('(') + len('('):final.rfind(',')]
                y = final[final.find(',') + len(','):final.rfind(')')]
                add_square(0.1, 50.0, 50.0, x, y, 0.5, len(squares))

        # ticks_to_next_ball -= 1
        # if ticks_to_next_ball <= 0:
        #     ticks_to_next_ball = 100
        #     mass = 0.1
        #     radius = 25
        #     inertia = pymunk.moment_for_circle(mass, 0, radius, (0, 0))
        #     body = pymunk.Body(mass, inertia)
        #     pos = pygame.mouse.get_pos()
        #     final = str(pos)
        #     x = final[final.find('(')+len('('):final.rfind(',')]
        #     print(x)
        #     y = final[final.find(',') + len(','):final.rfind(')')]
        #     print(int(x), int(y))
        #     body.position = int(x), (600 - int(y))
        #     shape = pymunk.Circle(body, radius, (0, 0))
        #     shape.friction = 0.5
        #     space.add(body, shape)
        #     balls.append(shape)

        ### Clear screen
        background = (160, 194, 205,)
        screen.fill(background) #Dull blueish color

        ### Draw stuff
        space.debug_draw(draw_options)
        rect = Rect(450, 50, 100, 50)
        pygame.draw.rect(screen, THECOLORS["navy"], (450, 50, 100, 50))

        #Draws text onto the rectangle GO button
        def text_objects(text, font):
            textSurface = font.render(text, True, THECOLORS["white"])
            return textSurface, textSurface.get_rect()

        smallText = pygame.font.Font("freesansbold.ttf", 20)
        textSurf, textRect = text_objects("GO!", smallText)
        textRect.center = ((450 + (100 / 2)), (50 + (50 / 2)))
        screen.blit(textSurf, textRect)

        ### Update physics
        dt = 1.0 / 60.0
        for x in range(1):
            space.step(dt)

        ### Flip screen
        pygame.display.flip()
        clock.tick(50)
        #pygame.display.set_caption("fps: " + str(clock.get_fps()))
        pygame.display.set_caption("ATWOOD SIMULATION")

#adds a circle to the screen with changeabe mass, radius, x and y position, and friction coef
def add_circle(mass, radius, xpos, ypos, friction, index):
    inertia = pymunk.moment_for_circle(mass, 0, radius, (0, 0))
    body = pymunk.Body(mass*5, inertia)
    body.position = int(xpos), (600 - int(ypos))
    shape = pymunk.Circle(body, radius, (0, 0))
    shape.color = pygame.color.THECOLORS["black"]
    shape.friction = friction
    space.add(body, shape)
    balls.insert(index, shape)

# adds a square to the screen with changeable mass, width, height, x and y positions, and friction coef
def add_square(mass, width, height, xpos, ypos, friction, index):
    inertia = pymunk.moment_for_box(mass, (width, height))
    body = pymunk.Body(mass*5, inertia)
    body.position = int(xpos), (600 - int(ypos))
    shape = pymunk.Poly.create_box(body, (width, height), 0)  # adding a radius (third param) bevels corners of poly
    shape.color = pygame.color.THECOLORS["black"]
    shape.friction = friction
    space.add(body, shape)

    squares.insert(index, shape)

# adds a right triangle to the screen with changeable mass, degree, x and y position, and friction coef
def add_ramp(mass, degree, xpos, ypos, friction, index):
    tan = math.tan((degree * (math.pi / 180)))
    height = 100 / tan  #height adjusts to make degree applicable
    inertia = pymunk.moment_for_poly(pymunk.inf, [(0, height), (100, 0), (0, 0)], (0, 0), 0)  # the length is always 100
    body = pymunk.Body(mass, inertia)
    body.position = int(xpos), (600 - int(ypos))
    shape = pymunk.Poly(body, [(0, height), (100, 0), (0, 0)])  # adding a radius (a third param) bevels corners of poly
    shape.color = pygame.color.THECOLORS["black"]
    shape.friction = friction
    space.add(body, shape)

    ramps.insert(index, shape)

# adds a right triangle to the screen with changeable mass, degree, x and y position, and friction coef
def add_ramp_flipped(mass, degree, xpos, ypos, friction, index):
    tan = math.tan((degree * (math.pi / 180)))
    height = 100 / tan  #height adjusts to make degree applicable
    inertia = pymunk.moment_for_poly(pymunk.inf, [(0, height), (-100, 0), (0, 0)], (0, 0), 0)  # the length is always 100
    body = pymunk.Body(mass, inertia)
    body.position = int(xpos), (600 - int(ypos))
    shape = pymunk.Poly(body, [(0, height), (-100, 0), (0, 0)])  # adding a radius (a third param) bevels corners of poly
    shape.color = pygame.color.THECOLORS["black"]
    shape.friction = friction
    space.add(body, shape)

    flipped_ramps.insert(index, shape)

if __name__ == '__main__':
    sys.exit(main())


from nisnap.utils.aseg import basal_ganglia_labels, cortical_labels
from nisnap.utils.aseg import amygdala_nuclei, hippocampal_subfields

__format__ = '.png'

def pick_labels(data, labels):
    """Replaces all values by zero in a Numpy array (`data`) except values
    passed in `labels`. Used to select some specific labels in a given
    segmentation."""

    import numpy as np
    data2 = np.where(np.isin(data, labels), data, 0)
    return data2

def _aget_cmap_(labels=[]):
    import json
    import numpy as np
    import nisnap
    import os.path as op
    n_labels = len(labels)
    fp = op.join(op.dirname(nisnap.__file__), 'utils', 'colormap.json')
    LUT = json.load(open(fp))
    LUT = {int(k): v for k, v in list(LUT.items())}

    if n_labels is None:
        n_labels = len(list(LUT.keys()))
    max_label = int(np.max(labels))
    LUT = [LUT.get(i, [0, 0, 0]) for i in range(0, max_label + 1)]
    LUT = np.array(LUT)
    LUT = LUT / 255.0
    return LUT

def _plot_contours_in_slice_(slice_seg, target_axis, labels=None):
    """Plots contour around the data in slice (after binarization)"""
    import numpy as np

    if labels is None:  # if n_labels is not provided then take max from slice
        labels = list(np.unique(slice_seg))

    from matplotlib.colors import ListedColormap
    cmap = ListedColormap(_aget_cmap_(labels))

    num_labels = len(cmap.colors)
    unique_labels = np.arange(num_labels, dtype='int16')

    from matplotlib import cm, colors
    normalize_labels = colors.Normalize(vmin=0, vmax=num_labels, clip=True)
    seg_mapper = cm.ScalarMappable(norm=normalize_labels, cmap=cmap)
    unique_labels_display = np.array(unique_labels)  # np.setdiff1d(unique_labels, 0)

    color_for_label = seg_mapper.to_rgba(unique_labels_display)

    from matplotlib import pyplot as plt
    plt.sca(target_axis)

    for index, label in enumerate(unique_labels_display):
        binary_slice_seg = slice_seg == index
        if not binary_slice_seg.any():
            continue
        plt.contour(binary_slice_seg,
                    levels=[0.5, ],
                    colors=(color_for_label[index],),
                    linewidths=1,
                    alpha=1,
                    zorder=1)

    return

def _snap_contours_(data, slices, axis, bg, figsize=None, bb=None, pbar=None):
    from matplotlib import pyplot as plt
    import numpy as np
    import tempfile
    import os

    plt.style.use('dark_background')

    paths = []
    fd, path = tempfile.mkstemp(suffix='_%s%s' % (axis, __format__))
    os.close(fd)
    paths.append(path)

    same_box = bb is not None
    if not same_box:
        bb = {}

    lambdas = {'x': lambda y, x: y[:, :, x],
               'y': lambda y, x: y[:, x, :],
               'z': lambda y, x: y[x, :, :]}

    labels = list(np.unique(data))

    if figsize is None:
        ratio = len(slices) / float(len(slices[0]))
        figsize = (figsize, figsize * ratio)

    fig = plt.figure(dpi=300, figsize=figsize)

    abs_index = 0
    for a, chunk in enumerate(slices):
        if not same_box:
            bb[a] = []

        for i, slice_index in enumerate(chunk):
            abs_index += 1

            ax = fig.add_subplot(len(slices), len(slices[0]), abs_index,
                                 label='%s_%s' % (axis, slice_index))
            test = np.flip(np.swapaxes(np.abs(lambdas[axis](data,
                                                            int(slice_index))),
                                       0, 1), 0)
            if not same_box:
                xs, ys = np.where(test != 0)
                bb[a].append((xs, ys))
            else:
                xs, ys = bb[a][i]

            if len(xs) == 0:
                continue

            test3 = np.flip(np.swapaxes(np.abs(lambdas[axis](bg,
                                                             int(slice_index))),
                                        0, 1), 0)
            test3 = test3[min(xs):max(xs) + 1, min(ys):max(ys) + 1]

            ax.imshow(test3, interpolation='none', cmap='gray')

            test = test[min(xs):max(xs) + 1, min(ys):max(ys) + 1]

            _plot_contours_in_slice_(test, ax, labels=labels)
            ax.axis('off')

            ax.text(0, 0, '%i' % slice_index,
                    {'color': 'w', 'fontsize': 10}, va="bottom", ha="left")

            if pbar is not None:
                pbar.update(1)

    fig.savefig(path, facecolor=fig.get_facecolor(),
                bbox_inches='tight',
                transparent=True,
                pad_inches=0)
    return paths, bb

def _snap_slices_(data, slices, axis, bb=None, figsize=None, pbar=None):
    from matplotlib import pyplot as plt
    import numpy as np
    import tempfile
    import os

    labels = list(np.unique(data))
    has_bb = bb is not None
    d = data.ravel()
    ratio = len(d[d <= 10]) / len(d)
    # has_orig = len(labels) > 50  # not bb is None
    is_raw = ratio < 0.80

    paths = []
    if not has_bb:
        bb = {}

    fig = plt.figure(dpi=300, figsize=figsize)

    from nisnap.utils.slices import __get_lambdas__
    lambdas = __get_lambdas__(data)

    fd, path = tempfile.mkstemp(suffix='_%s%s' % (axis, __format__))
    os.close(fd)
    paths.append(path)

    abs_index = 0
    for a, chunk in enumerate(slices):

        if not has_bb:
            bb[a] = []

        for i, slice_index in enumerate(chunk):
            abs_index += 1
            ax = fig.add_subplot(len(slices), len(slices[0]), abs_index,
                                 label='%s_%s' % (axis, slice_index))

            test = np.flip(np.swapaxes(np.abs(lambdas[axis](int(slice_index))),
                                       0, 1), 0)
            if not has_bb:
                xs, ys = list(np.where(test != 0))[:2]
                bb[a].append((xs, ys))
            else:
                xs, ys = bb[a][i]

            if len(xs) == 0:
                continue

            if len(data.shape) == 4:  # RGB mode (4D volume)
                test = test[min(xs):max(xs) + 1, min(ys):max(ys) + 1, :]
                ax.imshow((test * 255).astype(np.uint8), interpolation='none',)

            else:  # standard 3D label volume

                if is_raw:
                    vmax, cmap = (None, 'gray')
                else:
                    vmax = np.max(labels)
                    from matplotlib.colors import ListedColormap
                    cmap = ListedColormap(_aget_cmap_(labels))

                test = test[min(xs):max(xs) + 1, min(ys):max(ys) + 1]

                ax.imshow(test, interpolation='none', cmap=cmap,
                          vmin=0, vmax=vmax)

            ax.axis('off')
            ax.text(0, 0, '%i' % slice_index,
                    {'color': 'w', 'fontsize': 10}, va="bottom", ha="left")

            if pbar is not None:
                pbar.update(1)

    fig.savefig(path, facecolor=fig.get_facecolor(),
                bbox_inches='tight', transparent=True, pad_inches=0)
    return paths, bb

def __snap__(data, axes='xyz', bg=None, slices=None, rowsize=None,
             contours=False, figsize=None, samebox=False, margin=5):
    from matplotlib import pyplot as plt
    import logging as log

    plt.rcParams['figure.facecolor'] = 'black'
    plt.rcParams.update({'figure.max_open_warning': 0})

    from nisnap.utils.slices import cut_slices, _fix_rowsize_, _fix_figsize_
    from nisnap.utils.slices import __maxsize__

    rowsize = _fix_rowsize_(axes, rowsize)
    figsize = _fix_figsize_(axes, figsize)

    t = int(__maxsize__(data)/3.0)
    slices = cut_slices(data, axes, slices=slices, rowsize=rowsize,
                        threshold=t)
    n_slices = sum([sum([len(each) for each in slices[e]]) for e in axes])
    if n_slices == 0:
        msg = 'Should provide at least one slice. %s' % slices
        raise Exception(msg)
    has_orig = bg is not None

    if has_orig:
        n_slices = 2 * n_slices

    from tqdm import tqdm
    pbar = tqdm(total=n_slices, leave=False)

    paths, paths_orig = {}, {}

    for axis in axes:
        if samebox:
            from nisnap.utils.slices import __get_abs_minmax
            same_bb = __get_abs_minmax(data, axis, slices[axis], margin=margin)
            log.warning('Using bounding box: %s (axis %s)'
                        % (same_bb[0][0], axis))

        opt = {'slices': slices[axis],
               'axis': axis,
               'figsize': figsize[axis],
               'pbar': pbar}

        if contours:
            # Rendering contours
            if samebox:
                opt['bb'] = same_bb

            path, bb = _snap_contours_(data, bg=bg, **opt)
            paths[axis] = path

        else:
            opt['bb'] = None if not samebox else same_bb
            # Rendering masks
            path, bb = _snap_slices_(data, **opt)
            paths[axis] = path

        if has_orig:
            opt['bb'] = bb if not samebox else same_bb
            path, _ = _snap_slices_(bg, **opt)
            paths_orig[axis] = path

    pbar.update(n_slices)
    pbar.close()
    return paths, paths_orig

def __stack_img__(filepaths):
    import nibabel as nib
    import numpy as np

    stack = [np.asarray(nib.load(e).dataobj) for e in filepaths[:]]

    data = np.stack(stack, axis=-1)

    data2 = np.argmax(data, axis=-1) + 1
    black_pixels_mask = np.all(data == [0, 0, 0], axis=-1)
    data2[black_pixels_mask] = 0
    return data2

def plot_segment(filepaths, axes='xyz', bg=None, opacity=90, slices=None,
                 animated=False, savefig=None, contours=False, rowsize=None,
                 figsize=None, samebox=False, labels=None, margin=5):
    """Plots a set of segmentation maps/masks.

    Parameters
    ----------
    filepaths: a list of str
        Paths to segmentation maps (between 1 and 3). Must be of same
        dimensions and in same reference space.

    axes: string, or a tuple of str
        Choose the direction of the cuts (among 'x', 'y', 'z')

    bg: None or str
        Path to the background image that the masks will be plotted on top of.
        If nothing is specified, the segmentation maps/masks will be plotted
        only. The opacity (in %) of the segmentation maps when plotted over a
        background image. Only used if a background image is provided.
        Default: 10

    slices: None, or a tuple of floats
        The indexes of the slices that will be rendered. If None is given, the
        slices are selected automatically.

    animated: boolean, optional
        If True, the snapshot will be rendered as an animated GIF.
        If False, the snapshot will be rendered as a static PNG image. Default:
        False

    savefig: string, optional
        Filepath where the resulting snapshot will be created. If None is
        given, a temporary file will be created and/or the result will be
        displayed inline in a Jupyter Notebook.

    contours: boolean, optional
        If True, segmentations will be rendered as contoured regions. If False,
        will be rendered as superimposed masks. Default: False

    rowsize: None, or int, or dict
        Set the number of slices per row in the final compiled figure.
        Default: {'x': 9, 'y': 9, 'z': 6}

    figsize: None, or float
        Figure size (in inches) (matplotlib definition). Default: auto

    samebox: boolean, optional
        If True, bounding box will be fixed. If False, adjusted for each slice.

    labels: None or a tuple of int
        If a list of labels is provided, the label volume will be filtered to
        keep these labels only and remove the others.
        (Works with a label volume only, not with RGB mode)

    See Also
    --------
    xnat.plot_segment : To plot segmentation maps directly providing their
        experiment_id on an XNAT instance
    """
    import matplotlib
    import numpy as np
    import logging as log
    import os
    import tempfile
    matplotlib.use('Agg')

    fp = savefig
    if savefig is None:
        if animated:
            f, fp = tempfile.mkstemp(suffix='.gif')
        else:
            f, fp = tempfile.mkstemp(suffix=__format__)
        os.close(f)

    from nisnap.utils.parse import __check_axes__
    axes = __check_axes__(axes)

    # Creating snapshots (along given axes and original if needed)
    log.info('* Creating snapshots...')

    # Loading images
    import nibabel as nib
    if isinstance(filepaths, list):  # RGB mode
        data = __stack_img__(filepaths)
        log.info('* RGB mode')

    elif isinstance(filepaths, str):  # 3D label volume
        log.info('* Label volume')
        if labels is not None:
            from nisnap.utils import aseg
            filepaths = aseg.__picklabel_fs__(filepaths, labels=labels)
        data = np.asarray(nib.load(filepaths).dataobj)

    if bg is not None:
        bg = np.asarray(nib.load(bg).dataobj)

    paths, paths_orig = __snap__(data, axes=axes, bg=bg,
                                 slices=slices, contours=contours,
                                 rowsize=rowsize, figsize=figsize,
                                 samebox=samebox, margin=margin)

    from nisnap.utils.montage import __montage__
    has_orig = bg is not None
    __montage__(paths, paths_orig, axes, opacity, has_orig, animated,
                savefig=fp)

    if savefig is None:
        # Return image
        from IPython.display import Image
        return Image(filename=fp)


import numpy as np
import matplotlib.pyplot as plt

def read_jam_factor(path):
    last_time = 0
    data = []
    with open(path) as f:
        line = f.readline()
        while line:
            s = line.split(' ')
            s = [float(x) for x in s]
            t = s[0]

            if int(t) >= last_time + 30:
                data.append((t,
                    s[4] / s[1],
                    s[5] / s[2],
                    s[6] / s[3]))
                print(path)
                print(sum(s[4:7]) / sum(s[1:4]))
                last_time = int(t)
            line = f.readline()
    return data

data_gamma = read_jam_factor('statistics.gamma.log')
data_simple = read_jam_factor('statistics.ttc.log')

avg_total = 0
avg_count = 0
for x in data_gamma:
    avg_total += x[1] + x[2] + x[3]

fig, axs = plt.subplots(figsize=(17.5, 8))
axs.set_xlabel('Simulation Time (s)', fontsize=34)
axs.set_ylabel('Congestion Factor', fontsize=34)
plt.xlim(0, 1200)
plt.ylim(-0.01, 1.19)
plt.rc('legend',fontsize=26)
axs.tick_params(axis='x', labelsize=32)
axs.tick_params(axis='y', labelsize=32)

# GAMMA Car Jam Factor.
axs.plot(
        [x[0] for x in data_gamma],
        [x[1] for x in data_gamma],
        label='Car (C-GAMMA)',
        color='red',
        linestyle='solid',
        linewidth=3)

# TTC Car Jam Factor.
axs.plot(
        [x[0] for x in data_simple],
        [x[1] for x in data_simple],
        label='Car (TTC)',
        color='red',
        linestyle=(0, (5, 5)),
        linewidth=3)

# GAMMA Bike Jam Factor.
axs.plot(
        [x[0] for x in data_gamma],
        [x[2] for x in data_gamma],
        label='Bike (C-GAMMA)',
        color='green',
        linestyle='solid',
        linewidth=3)

# TTC Bike Jam Factor.
axs.plot(
        [x[0] for x in data_simple],
        [x[2] for x in data_simple],
        label='Bike (TTC)',
        color='green',
        linestyle=(0, (5, 5)),
        linewidth=3)

# GAMMA Ped. Jam Factor.
axs.plot(
        [x[0] for x in data_gamma],
        [x[3] for x in data_gamma],
        label='Pedestrian (C-GAMMA)',
        color='blue',
        linestyle='solid',
        linewidth=3)

# TTC Ped. Jam Factor.
axs.plot(
        [x[0] for x in data_simple],
        [x[3] for x in data_simple],
        label='Pedestrian (TTC)',
        color='blue',
        linestyle=(0, (5, 5)),
        linewidth=3)

axs.legend(loc="upper right", ncol=3)

fig.tight_layout()
plt.show()