#!/usr/bin/env python3

from lxml import etree as ET
import base64
from . computations import Computations
import sys

nsmap = {
    "http://www.hmrc.gov.uk/schemas/ct/comp/2023-01-01": "ct-comp",
    "http://www.hmrc.gov.uk/schemas/ct/dpl/2021-01-01": "dpl",
    "http://xbrl.frc.org.uk/fr/2023-01-01/core": "uk-core"
}

ct_ns = "http://www.govtalk.gov.uk/taxation/CT/5"

class Box:
    def __init__(self, id, kind=None):
        self.kind = kind
        self.id = id

    def present(self, obj):
        if self.id in obj.form_values["ct600"]:

            value = obj.form_values["ct600"][self.id]
            if value == None:
                return False

            return True

        return False

    def fetch(self, obj):
        return obj.form_values["ct600"][self.id]

    def get(self, obj):
        value = self.fetch(obj)

        if self.kind == "yesno":
            return "yes" if value else "no"

        if self.kind == "money":
            return "%.2f" % value

        if self.kind == "rate":
            return "%.2f" % value

        if self.kind == "pounds":
            return "%.2f" % int(value)

        if self.kind == "yes":
            return "yes" if value else "FIXME"

        if self.kind == "date":
            return str(value)

        if self.kind == "year":
            return str(value)

        if self.kind == "companytype":
            return "%02d" % value

        return value

class Fixed(Box):
    def __init__(self, value, kind=None):
        self.value = value
        self.kind = kind

    def present(self, obj):
        return True

    def fetch(self, obj):
        return self.value

class InputBundle:

    def __init__(self, comps, accts, form_values, params, atts):
        self.comps = comps
        self.accts = accts
        self.form_values = form_values
        self.params = params
        self.atts = atts

    def box(self, num):
        val = self.form_values["ct600"][num]
        if val is None: return ""
        return val

    def date(self, num):
        val = self.form_values["ct600"][num]
        if val is None: return ""
        return str(val)

    def irheader(self):

        irheader = ET.Element(
            "{%s}%s" % (ct_ns, "IRheader"),
            nsmap={"ct": ct_ns}
        )

        keys = ET.SubElement(irheader, "{%s}%s" % (ct_ns, "Keys"))

        ET.SubElement(
            keys, "{%s}%s" % (ct_ns, "Key"), Type="UTR"
        ).text = str(self.box(3))

        ET.SubElement(
            irheader, "{%s}%s" % (ct_ns, "PeriodEnd")
        ).text = str(self.date(35))

        principal = ET.SubElement(irheader, "{%s}%s" % (ct_ns, "Principal"))
        contact = ET.SubElement(principal, "{%s}%s" % (ct_ns, "Contact"))
        name = ET.SubElement(contact, "{%s}%s" % (ct_ns, "Name"))
        ET.SubElement(
            name, "{%s}%s" % (ct_ns, "Ttl")
        ).text = self.params.get("title")
        ET.SubElement(
            name, "{%s}%s" % (ct_ns, "Fore")
        ).text = self.params.get("first-name")
        ET.SubElement(
            name, "{%s}%s" % (ct_ns, "Sur")
        ).text = self.params.get("second-name")
        ET.SubElement(
            contact, "{%s}%s" % (ct_ns, "Email")
        ).text = self.params.get("email")
        tel = ET.SubElement(contact, "{%s}%s" % (ct_ns, "Telephone"))
        ET.SubElement(
            tel, "{%s}%s" % (ct_ns, "Number")
        ).text = self.params.get("phone")

        ET.SubElement(irheader, "{%s}%s" % (ct_ns, "IRmark")).text = ""
        ET.SubElement(irheader, "{%s}%s" % (ct_ns, "Sender")).text = "Company"

        return irheader

    def get_return(self):

        comp_ixbrl = base64.b64encode(self.comps).decode("utf-8")
        accounts_ixbrl = base64.b64encode(self.accts).decode("utf-8")

        mapping = {
            "CompanyInformation": {
                "CompanyName": Box(1),
                "RegistrationNumber": Box(2),
                "Reference": Box(3),
                "CompanyType": Box(4, kind="companytype"),
                "NorthernIreland": {
                    "NItradingActivity": Box(5, kind="yesno"),
                    "SME": Box(6, kind="yesno"),
                    "NIemployer": Box(7, kind="yesno"),
                    "SpecialCircumstances": Box(8, kind="yesno"),
                },
                "PeriodCovered": {
                    "From": Box(30, kind="date"),
                    "To": Box(35, kind="date"),
                }
            },
            "ReturnInfoSummary": {
                "ThisPeriod": Box(40, kind="yesno"),
                "EarlierPeriod": Box(45, kind="yesno"),
                "MultipleReturns": Box(50, kind="yesno"),
                "ProvisionalFigures": Box(55, kind="yesno"),
                "PartOfNonSmallGroup": Box(60, kind="yesno"),
                "RegisteredAvoidanceScheme": Box(65, kind="yesno"),
                "TransferPricing": {
                    "Adjustment": Box(70, kind="yesno"),
                    "SME": Box(75, kind="yesno"),
                },
                "Accounts": {
                    "ThisPeriodAccounts": Box(80, kind="yesno"),
                    "DifferentPeriod": Box(85, kind="yesno"),
                    "NoAccountsReason": Box(90),
                },
                "Computations": {
                    "ThisPeriodComputations": Box(80, kind="yesno"),
                    "DifferentPeriod": Box(85, kind="yesno"),
                    "NoComputationsReason": Box(90),
                },
                "SupplementaryPages": {
                    "CT600A": Box(95, kind="yesno"),
                    "CT600B": Box(100, kind="yesno"),
                    "CT600C": Box(105, kind="yesno"),
                    "CT600D": Box(110, kind="yesno"),
                    "CT600E": Box(115, kind="yesno"),
                    "CT600F": Box(120, kind="yesno"),
                    "CT600G": Box(125, kind="yesno"),
                    "CT600H": Box(130, kind="yesno"),
                    "CT600I": Box(135, kind="yesno"),
                    "CT600J": Box(140, kind="yesno"),
                    "CT600K": Box(141, kind="yesno"),
                    "CT600L": Box(142, kind="yesno"),
                    "CT600N": Box(144, kind="yesno"),
                }
            },
            "Turnover": {
                "Total": Box(145, kind="pounds"),
                "OtherFinancialConcerns": Box(150, kind="pounds"),
            },
            "CompanyTaxCalculation": {
                "Income": {
                    "Trading": {
                        "Profits": Box(155, kind="pounds"),
                        "LossesBroughtForward": Box(160, kind="pounds"),
                        "NetProfits": Box(165, kind="pounds"),
                    },
                    "NonTradingLoanProfitsAndGains": Box(170, kind="pounds"),
                    "IncomeStatedNet": Box(172, kind="yesno"),
                    "NonLoanAnnuitiesAnnualPaymentsDiscounts": Box(175, kind="pounds"),
                    "NonUKdividends": Box(180, kind="pounds"),
                    "DeductedIncome": Box(185, kind="pounds"),
                    "PropertyBusinessIncome": Box(190, kind="pounds"),
                    "NonTradingGainsIntangibles": Box(195, kind="pounds"),
                    "TonnageTaxProfits": Box(200, kind="pounds"),
                    "OtherIncome": Box(205, kind="pounds"),
                },
                "ChargeableGains": {
                    "GrossGains": Box(210, kind="pounds"),
                    "AllowableLosses": Box(215, kind="pounds"),
                    "NetChargeableGains": Box(220, kind="pounds"),
                },
                "LossesBroughtForward": Box(225, kind="pounds"),
                "NonTradeDeficitsOnLoans": Box(230, kind="pounds"),
                "ProfitsBeforeOtherDeductions": Box(235, kind="pounds"),
                "DeductionsAndReliefs": {
                    "UnquotedShares": Box(240, kind="pounds"),
                    "ManagementExpenses": Box(245, kind="pounds"),
                    "UKpropertyBusinessLosses": Box(250, kind="pounds"),
                    "CapitalAllowances": Box(255, kind="pounds"),
                    "NonTradeDeficits": Box(260, kind="pounds"),
                    "CarriedForwardNonTradeDeficits": Box(263, kind="pounds"),
                    "NonTradingLossIntangibles": Box(265, kind="pounds"),
                    "TradingLosses": Box(275, kind="pounds"),
                    "TradingLossesCarriedBack": Box(280, kind="yesno"),
                    "TradingLossesCarriedForward": Box(285, kind="pounds"),
                    "NonTradeCapitalAllowances": Box(290, kind="pounds"),
                    "Total": Box(295, kind="pounds"),
                },
                "ChargesAndReliefs": {
                    "ProfitsBeforeDonationsAndGroupRelief": Box(
                        300, kind="pounds"
                    ),
                    "QualifyingDonations": Box(305, kind="pounds"),
                    "GroupRelief": Box(310, kind="pounds"),
                    "GroupReliefForCarriedForwardLosses": Box(312, kind="pounds"),
                },
                "ChargeableProfits": Box(315, kind="pounds"),
                "RingFenceProfitsIncluded": Box(320, kind="pounds"),
                "NorthernIrelandProfitsIncluded": Box(325, kind="pounds"),
                "CorporationTaxChargeable": {
                    "AssociatedCompanies": {
                        "ThisPeriod": Box(326),
                        "AssociatedCompaniesFinancialYears": {
                            "FirstYear": Box(327),
                            "SecondYear": Box(328)
                        },
                        "StartingOrSmallCompaniesRate": Box(329, kind="yesno")
                    },
                    "FinancialYearOne": {
                        "Year": Box(330, kind="year"),
                        "Details": {
                            "Profit": Box(335, kind="pounds"),
                            "TaxRate": Box(340, kind="rate"),
                            "Tax": Box(345, kind="money")
                        }
                    },
                    "FinancialYearTwo": {
                        "Year": Box(380, kind="year"),
                        "Details": {
                            "Profit": Box(385, kind="pounds"),
                            "TaxRate": Box(390, kind="rate"),
                            "Tax": Box(395, kind="money")
                        }
                    },
                },
                "CorporationTax": Box(430, kind="money"),
                "MarginalReliefForRingFenceTrades": Box(435, kind="money"),
                "NetCorporationTaxChargeable": Box(440, kind="money"),
                "TaxReliefsAndDeductions": {
                    "CommunityInvestmentRelief": Box(445, kind="money"),
                    "DoubleTaxation": {
                        "DoubleTaxationRelief": Box(450, kind="money"),
                        "UnderlyingRate": Box(455, kind="yesno"),
                        "AmountCarriedBack": Box(460, kind="yes"),
                    },
                    "AdvancedCorporationTax": Box(465, kind="money"),
                    "TotalReliefsAndDeductions": Box(470, kind="money"),
                },
                "CJRS": {
                    "CJRSreceived": Box(471, kind="money"),
                    "CJRSdue": Box(472, kind="money"),
                    "CJRSoverpaymentAlreadyAssessed": Box(473, kind="money"),
                    "JobRetentionBonusOverpayment": Box(474, kind="money"),
                },
            },
            "CalculationOfTaxOutstandingOrOverpaid": {
                "NetCorporationTaxLiability": Box(475, kind="money"),
                "LoansToParticipators": Box(480, kind="money"),
                "CT600AreliefDue": Box(485, kind="yesno"),
                "CFCtaxPayable": Box(490, kind="money"),
                "BankLevyPayable": Box(495, kind="money"),
                "BankSurchargePayable": Box(496, kind="money"),
                "RPDTpayable": Box(497, kind="money"),
                "CFCandBankLevyTotal": Box(500, kind="money"),
                "EOGPLpayable": Box(501, kind="money"),
                "SupplementaryCharge": Box(505, kind="money"),
                "TaxChargeable": Box(510, kind="money"),
                "IncomeTax": {
                    "DeductedIncomeTax": Box(515, kind="money"),
                    "TaxRepayable": Box(520, kind="money"),
                },
                "TaxPayable": Box(525, kind="money"),
                "CJRSoverpaymentsNowDue": Box(526, kind="money"),
                "RestitutionTax": Box(527, kind="money"),
                "TaxPayableIncludingRestitutionTax": Box(528, kind="money"),
            },
            "TaxReconciliation": {
                "ResearchAndDevelopmentCredit": Box(530, kind="money"),
                "VaccineCredit": Box(535, kind="money"),
                "CreativeCredit": Box(540, kind="money"),
                "ResearchAndDevelopmentVaccineOrCreativeTaxCredit":  Box(545, kind="money"),
                "LandRemediationCredit": Box(550, kind="money"),
                "LifeAssuranceCompanyCredit": Box(555, kind="money"),
                "LandOrLifeCredit": Box(560, kind="money"),
                "CapitalAllowancesFirstYearCredit": Box(565, kind="money"),
                "SurplusResearchAndDevelopmentCreditsOrCreativeCreditPayable": Box(570, kind="money"),
                "LandOrLifeCreditPayable": Box(575, kind="money"),
                "CapitalAllowancesFirstYearCreditPayable": Box(580, kind="money"),
                "RingFenceCorpTaxIncluded": Box(585, kind="money"),
                "NIcorporationTaxIncluded": Box(586, kind="money"),
                "RingFenceSupplementaryChargeIncluded": Box(590, kind="money"),
                "TaxAlreadyPaid": Box(595, kind="money"),
                "TaxOutstandingOrOverpaid": {
                    "TaxOutstanding": Box(600, kind="money"),
                    "TaxOverpaid": Box(605, kind="money"),
                },
                "RefundsSurrendered": Box(610, kind="money"),
                "RandDExpenditureCreditsSurrendered": Box(615, kind="money"),
            },
            "IndicatorsAndInformation": {
                "FrankedInvestmentIncome": Box(620, kind="pounds"),
                "NumberOf51groupCompanies": Box(625),
                "InstalmentPayments": Box(630, kind="yes"),
                "VeryLargeQIPs": Box(631, kind="yes"),
                "GroupPayment": Box(635, kind="yes"),
                "IntangibleAssets": Box(640, kind="yes"),
                "CrossBorderRoyalty": Box(645, kind="yes"),
                "EatOutToHelpOutScheme": Box(647, kind="pounds"),
            },
            "EnhancedExpenditure": {
                "SMEclaim": Box(650, kind="yes"),
                "LargeCompanyClaim": Box(655, kind="yes"),
                "RAndDClaimNotificationForm": Box(656, kind="yesno"),
                "AdditionalRAndDForm": Box(657, kind="yesno"),
                "RAndDExpenditureSME": Box(659, kind="pounds"),
                "RandDEnhancedExpenditure": Box(660, kind="pounds"),
                "CreativeEnhancedExpenditure": Box(665, kind="pounds"),
                "RandDAndCreativeEnhancedExpenditure": Box(670, kind="pounds"),
                "SMEclaimAsLargeCompany": Box(675, kind="pounds"),
                "VaccineResearch": Box(680, kind="pounds"),
            },
            "LandRemediationEnhancedExpenditure": Box(685, kind="pounds"),
            "AllowancesAndCharges": {
                "AIACapitalAllowancesInc": Box(690, kind="pounds"),
                "MachineryAndPlantSpecialRatePool": {
                    "BalancingCharges": Box(695, kind="pounds"),
                    "CapitalAllowances": Box(700, kind="pounds"),
                },
                "MachineryAndPlantMainPool": {
                    "BalancingCharges": Box(705, kind="pounds"),
                    "CapitalAllowances": Box(710, kind="pounds"),
                },
                "StructuresAndBuildingsCapitalAllowances": Box(711, kind="pounds"),
                "ElectricChargePoints": {
                    "BalancingCharges": Box(713, kind="pounds"),
                    "CapitalAllowances": Box(714, kind="pounds"),
                },
                "BusinessPremisesRenovationIncluded": {
                    "BalancingCharges": Box(715, kind="pounds"),
                    "CapitalAllowances": Box(720, kind="pounds"),
                },
                "EnterpriseZones": {
                    "BalancingCharges": Box(721, kind="pounds"),
                    "CapitalAllowances": Box(722, kind="pounds"),
                },
                "ZeroEmissionsGoodsVehicles": {
                    "BalancingCharges": Box(723, kind="pounds"),
                    "CapitalAllowances": Box(724, kind="pounds"),
                },
                "ZeroEmissionsCars": {
                    "BalancingCharges": Box(726, kind="pounds"),
                    "CapitalAllowances": Box(727, kind="pounds"),
                },
                "Others": {
                    "BalancingCharges": Box(725, kind="pounds"),
                    "CapitalAllowances": Box(730, kind="pounds"),
                },
            },
            "NotIncluded": {
                "AIACapitalAllowancesNotInc": Box(735, kind="pounds"),
                "StructuresAndBuildingsCapitalAllowances": Box(736, kind="pounds"),
                "ElectricChargePoints": {
                    "BalancingCharges": Box(737, kind="pounds"),
                    "CapitalAllowances": Box(738, kind="pounds"),
                },
                "BusinessPremisesRenovationNotIncluded": {
                    "BalancingCharges": Box(740, kind="pounds"),
                    "CapitalAllowances": Box(745, kind="pounds"),
                },
                "EnterpriseZones": {
                    "BalancingCharges": Box(746, kind="pounds"),
                    "CapitalAllowances": Box(747, kind="pounds"),
                },
                "ZeroEmissionsGoodsVehicles": {
                    "BalancingCharges": Box(748, kind="pounds"),
                    "CapitalAllowances": Box(749, kind="pounds"),
                },
                "ZeroEmissionsCars": {
                    "BalancingCharges": Box(751, kind="pounds"),
                    "CapitalAllowances": Box(752, kind="pounds"),
                },
                "OtherAllowancesAndCharges": {
                    "BalancingCharges": Box(750, kind="pounds"),
                    "CapitalAllowances": Box(755, kind="pounds"),
                },
            },
            "QualifyingExpenditure": {
                "MachineryAndPlantExpenditure": Box(760, kind="pounds"),
                "DesignatedEnvironmentallyFriendlyMachineryAndPlant": Box(765, kind="pounds"),
                "MachineryAndPlantLongLife": Box(770, kind="pounds"),
                "StructuresAndBuildings": Box(771, kind="pounds"),
                "OtherMachineryAndPlant": Box(775, kind="pounds"),
            },
            "LossesDeficitsAndExcess": {
                "AmountArising": {
                    "LossesOfTradesUK": {
                        "Arising": Box(780, kind="pounds"),
                        "SurrenderMaximum": Box(785, kind="pounds"),
                    },
                    "LossesOfTradesOutsideUK": Box(790, kind="pounds"),
                    "Loans": {
                        "Arising": Box(795, kind="pounds"),
                        "SurrenderMaximum": Box(800, kind="pounds"),
                    },
                    "UKpropertyBusinessLosses": {
                        "Arising": Box(805, kind="pounds"),
                        "SurrenderMaximum": Box(810, kind="pounds"),
                    },
                    "OverseasPropertyBusinessLosses": Box(815, kind="pounds"),
                    "MiscLosses": Box(820, kind="pounds"),
                    "CapitalLosses": Box(825, kind="pounds"),
                    "NonTradingLossesIntangibles": {
                        "Arising": Box(830, kind="pounds"),
                        "SurrenderMaximum": Box(835, kind="pounds"),
                    },
                },
                "ExcessAmounts": {
                    "NonTradeCapital": Box(840, kind="pounds"),
                    "QualifyingDonations": Box(845, kind="pounds"),
                    "ManagementExpenses": {
                        "Arising": Box(850, kind="pounds"),
                        "SurrenderMaximum": Box(855, kind="pounds"),
                    },
                }
            },
            "NorthernIrelandInformation": {
                "NIagainstUK": Box(856, kind="pounds"),
                "NIagainstNI": Box(857, kind="pounds"),
                "UKagainstNI": Box(858, kind="pounds"),
            },
            "OverpaymentsAndRepayments": {
                "OwnRepaymentsLowerLimit": Box(860, kind="pounds"),
                "RepaymentsForThePeriodCoveredByThisReturn": {
                    "CorporationTax": Box(865, kind="money"),
                    "IncomeTax": Box(870, kind="money"),
                    "RandDTaxCredit": Box(875, kind="money"),
                    "RandDExpenditureCredit": Box(880, kind="money"),
                    "CreativeCredit": Box(885, kind="money"),
                    "LandRemediationCredit": Box(890, kind="money"),
                    "PayableCapitalAllowancesFirstYearCredit": Box(895, kind="money"),
                },
                "Surrender": {
                    "Amount": Box(900, kind="money"),
                    "JointNotice": {
                        "Attached": Box(905, kind="yes"),
                        "WillFollow": Box(910, kind="yes"),
                    },
                    "StopUntilNotice": Box(915, kind="money"),
                },
                "BankAccountDetails": {
                    "BankName": Box(920),
                    "SortCode": Box(925),
                    "AccountNumber": Box(930),
                    "AccountName": Box(935),
                    "BuildingSocReference": Box(940),

                },
                "PaymentToPerson": {
                    "Recipient": Box(955),
                    "Address": {

                        # FIXME
                        "Line": Box(960),

                        # FIXME: AdditionalLine
                        # FIXME: PostCode

                    },
                    "NomineeReference": Box(965),
                },
            },
            "Declaration": {
                "AcceptDeclaration": Fixed("yes"),
                "Name": Box(975),
                "Status": Box(985),
            },
            "AttachedFiles": {
                "XBRLsubmission": {
                    "Computation": {
                        "Instance": {
                            "EncodedInlineXBRLDocument": Fixed(comp_ixbrl),
                        }
                    },
                    "Accounts": {
                        "Instance": {
                            "EncodedInlineXBRLDocument": Fixed(accounts_ixbrl),
                        }
                    }
                }
            },
            "EnergyProfitsLevy": {
                "EOGPLamounts": Box(986, kind="pounds"),
            }
        }

        mapping2 = {
            "CompanyInformation": {
                "CompanyName": Box(1),
                "RegistrationNumber": Box(2),
                "Reference": Box(3),
                "CompanyType": Box(4, kind="companytype"),
                "NorthernIreland": {
                    "NItradingActivity": Box(5, kind="yesno"),
                    "SME": Box(6, kind="yesno"),
                    "NIemployer": Box(7, kind="yesno"),
                    "SpecialCircumstances": Box(8, kind="yesno"),
                },
                "PeriodCovered": {
                    "From": Box(30, kind="date"),
                    "To": Box(35, kind="date"),
                }
            }
        }

        ctr = ET.Element(
            "{%s}%s" % (ct_ns, "CompanyTaxReturn"),
            nsmap={"ct": ct_ns}
        )

        def add_to_tree(root, tree):

            for key, value in tree.items():

                if type(value) == dict:

                    tag = "{%s}%s" % (ct_ns, key)
                    elt = ET.Element(tag)
                    add_to_tree(elt, value)
                    if len(elt) > 0:
                        root.append(elt)

                elif isinstance(value, Box):

                    if value.present(self):

                        tag = "{%s}%s" % (ct_ns, key)

                        eltval = value.get(self)

                        if isinstance(eltval, list):
                            for it in eltval:
                                itval = str(it)
                                ET.SubElement(root, tag).text = itval

                        else:
                            eltval = str(eltval)
                            ET.SubElement(root, tag).text = eltval

        add_to_tree(ctr, mapping)

        irenv = ET.Element(
            "{%s}%s" % (ct_ns, "IRenvelope"), nsmap={"ct": ct_ns}
        )

        irenv.append(self.irheader())
        irenv.append(ctr)

        # CompanyTaxReturn needs an attribute
        irenv.find(".//{%s}CompanyTaxReturn[1]" % ct_ns).set(
            "ReturnType", "new"
        )

        attchs = irenv.find(".//{%s}AttachedFiles[1]" % ct_ns)

        for f in self.atts:

            raw = self.atts[f]
            body = base64.b64encode(raw).decode("utf-8")

            att = ET.Element(
                "{%s}%s" % (ct_ns, "Attachment"),
                nsmap={"ct": ct_ns}
            )

            att.text = body

            att.set("Filename", f)
            att.set("Description", "supporting document")
            att.set("Format", "pdf")
            att.set("Size", str(len(raw)))
            att.set("Type", "other")

            attchs.append(att)

        return ET.ElementTree(irenv)


from tabulate import tabulate
import colorama
from colorama import Fore
import os

# Chess pieces are global to modify at run

board = {}

# Row 1
board['a1'] = 'R'
board['b1'] = 'N'
board['c1'] = 'B'
board['d1'] = 'Q'
board['e1'] = 'K'
board['f1'] = 'B'
board['g1'] = 'N'
board['h1'] = 'R'

# Row 2
board['a2'] = 'P'
board['b2'] = 'P'
board['c2'] = 'P'
board['d2'] = 'P'
board['e2'] = 'P'
board['f2'] = 'P'
board['g2'] = 'P'
board['h2'] = 'P'

# Row 3
board['a3'] = ' '
board['b3'] = ' '
board['c3'] = ' '
board['d3'] = ' '
board['e3'] = ' '
board['f3'] = ' '
board['g3'] = ' '
board['h3'] = ' '

# Row 4
board['a4'] = ' '
board['b4'] = ' '
board['c4'] = ' '
board['d4'] = ' '
board['e4'] = ' '
board['f4'] = ' '
board['g4'] = ' '
board['h4'] = ' '

# Row 5
board['a5'] = ' '
board['b5'] = ' '
board['c5'] = ' '
board['d5'] = ' '
board['e5'] = ' '
board['f5'] = ' '
board['g5'] = ' '
board['h5'] = ' '

# Row 6
board['a6'] = ' '
board['b6'] = ' '
board['c6'] = ' '
board['d6'] = ' '
board['e6'] = ' '
board['f6'] = ' '
board['g6'] = ' '
board['h6'] = ' '

# Row 7
board['a7'] = 'P'
board['b7'] = 'P'
board['c7'] = 'P'
board['d7'] = 'P'
board['e7'] = 'P'
board['f7'] = 'P'
board['g7'] = 'P'
board['h7'] = 'P'

# Row 8
board['a8'] = 'R'
board['b8'] = 'N'
board['c8'] = 'B'
board['d8'] = 'Q'
board['e8'] = 'K'
board['f8'] = 'B'
board['g8'] = 'N'
board['h8'] = 'R'

def drawBoard(move1, move2):

	if move1 != ' ':  # When move is empty print a clear start up position board
		clean_move1, piece1 = determinePiece(move1, 'W')
		clean_move2, piece2 = determinePiece(move2, 'B')

		board[clean_move1] = piece1
		board[clean_move2] = piece2

	table=[[board['a8'],board['b8'],board['c8'],board['d8'],board['e8'],board['f8'],board['g8'],board['h8']],
	     [board['a7'], board['b7'], board['c7'], board['d7'], board['e7'], board['f7'], board['g7'], board['h7']],
	     [board['a6'], board['b6'], board['c6'], board['d6'], board['e6'], board['f6'], board['g6'], board['h6']],
		 [board['a5'], board['b5'], board['c5'], board['d5'], board['e5'], board['f5'], board['g5'], board['h5']],
		 [board['a4'], board['b4'], board['c4'], board['d4'], board['e4'], board['f4'], board['g4'], board['h4']],
	     [board['a3'], board['b3'], board['c3'], board['d3'], board['e3'], board['f3'], board['g3'], board['h3']],
	     [board['a2'], board['b2'], board['c2'], board['d2'], board['e2'], board['f2'], board['g2'], board['h2']],
	     [board['a1'], board['b1'], board['c1'], board['d1'], board['e1'], board['f1'], board['g1'], board['h1']]]

	print(tabulate(table, tablefmt='fancy_grid'))

def determinePiece(move, color):
	keyPieces = ['R', 'N', 'B', 'Q', 'K']

	if move[0] not in keyPieces:
		piece = 'P'
		clean_move = move
		pawnMovement(clean_move, color)

	else:
		piece = move[0]
		clean_move = move[1:]

	return clean_move, piece

def pawnMovement(s, color):
	# two steps only if row 2 for white and 7 for Black everything else one to move forward
	s.split()  # get 1st the column then the number

	if color == 'W':
		# if row is 4 and row 2 is occupied decrease 2 movements
		x = s[0]
		y = int(s[1]) - 1 # while moves forward

	elif color == 'B':
		x = s[0]
		y = int(s[1]) + 1  # black  moves backwards

	clearCell = str(x) + str(y)  # while clears at the back minus
	board[clearCell] = ' '

# Main Function
drawBoard(' ',' ') #clear the board
drawBoard('d3', 'd6') # Draw pawn move
input("Press Enter to continue...")
os.system('cls||clear')
drawBoard('Nc3', 'Nc6') # Draw knight move
input("Press Enter to continue...")
os.system('cls||clear')
drawBoard('a3', 'h6')


import os

from dotenv import load_dotenv

from lof.env_vars import path_for_env_file, set_env_variables

current_dir = os.path.dirname(os.path.abspath(__file__))
test_json_vars = os.path.join(current_dir, "test_data/vars.json")
test_env_vars = os.path.join(current_dir, "test_data/.env")
lof_dir = ".lof"

def test_path_for_env_file():
    expected = ".lof/.env"
    os.makedirs(lof_dir, exist_ok=True)
    assert path_for_env_file(test_json_vars, lof_dir) == expected
    # check no issues
    assert load_dotenv(expected) is True

def test_set_env_variables_json():
    # check no issues
    assert set_env_variables(test_json_vars) is None
    assert os.environ["DB_USER"] == "test"

def test_set_env_variables_env():
    # check no issues
    assert set_env_variables(test_env_vars) is None
    assert os.environ["TEST_ENV"] == "DB_NAME_TEST"


def collatz(number):
    if number == 1:
        print( 'collatzs complete')
        restartCollatz()

    elif number % 2 == 0: # if even
        print (number // 2)
        collatz(number // 2)
    else:
        print (3 * number + 1) # if odd
        collatz(3* number +1)

def restartCollatz():
    print ( 'Do you want to choose another number?')
    print (' Y / N ')
    yesOrNo = str(input())
    if yesOrNo == 'Y':
        print('Type an interger in for Collatzs sequence')
        collatz(int(input()))
    else:
        print ( 'Thank you for using' )

print('Type an interger in for Collatzs sequence')

try:
    collatz(int(input()))
except ValueError:
    print( 'That wasnt an Interger was it?')
    print( 'Try again' )
    collatz(int(input()))


def next_bigger(n):
    # 1. 거꾸로
    n = str(n)
    data = [i for i in n ]

    if n == ''.join(sorted ( data , key= lambda x : -int(x) ) ) :
        return - 1

    # 2.
    cnt = 0
    for a , b in zip ( data[::-1] , data[-2::-1] ) :

        if a > b :
            print ( data , b, a , cnt )
            temp = data[ len(data) - cnt - 1:].copy()
            min_data = min ( list ( filter ( lambda x:x > b , temp ) ) )
            temp.remove(min_data)
            temp.append(b)
            temp = sorted( temp , key = lambda x : x)
            left = data[:len(data) - cnt - 2]
            left.append( min_data )
            return int ( ''.join(left + temp )  )
        cnt += 1

print ( next_bigger( 513 )     )

Next bigger number with the same digits.py


import json
from basics import Creature, KEYWORDS

# Load json (http://mtgjson.com/)
json_file = open("SOI.json")
set_json = json.load(json_file)["cards"]
for i in range(len(set_json)):
    for unwanted in ["artist", "flavor", "id", "imageName", "layout",
                     "multiverseid", "number", "type"]:
        try:
            set_json[i].pop(unwanted)
        except KeyError:
            continue
creature_jsons = [card for card in set_json if "Creature" in card["types"]]
for creature_json in creature_jsons:
    try:
        creature_json["keywords"] = [kw for kw in KEYWORDS if \
                                     kw in creature_json["text"]]
    except KeyError:
        continue
creature_list = [Creature(creature_json) for creature_json in creature_jsons]


import dataclasses
from typing import Dict

from neural_testbed import base
from neural_testbed.real_data import data_sampler
from neural_testbed.real_data import datasets

# problem_ids are strings of the form {sweep_name}{SEPARATOR}{index}.
SEPARATOR = '/'

@dataclasses.dataclass(frozen=True)
class ProblemConfig:
  """Problem configuration including prior knowledge and some hyperparams."""
  # Agent's a priori knowledge about the problem.
  prior_knowledge: base.PriorKnowledge
  # Random seed controlling all the randomness in the problem.
  seed: int
  # Name of the dataset
  dataset_name: str
  # Test sampler constructor
  test_sampler_ctor: data_sampler.SamplerCtor = data_sampler.make_local_sampler
  # Number of inputs (X's) used for evaluation.
  num_test_seeds: int = 1000
  # Number of samples generated from ENN during evaluation.
  num_enn_samples: int = 1000

  @property
  def meta_data(self):
    meta = dataclasses.asdict(self)
    meta.pop('prior_knowledge')
    meta.update(dataclasses.asdict(self.prior_knowledge))
    return meta

def classification_sweep(num_seed: int = 5,
                         initial_seed: int = 0,
                         temperature: float = 0.01) -> Dict[str, ProblemConfig]:
  """Generate hyperparameter sweep for classification problems.

  Args:
    num_seed: number of seeds per configuration of other hyperparameters.
    initial_seed: initial value of the seed.
    temperature: temperature to be used in prior_knowledge. It does not affect
      the data.

  Returns:
    Mapping problem_id: gp_settings (for use in gp_load).
  """
  configs = []
  for tau in [1, 10]:
    for dataset_name in datasets.CLASSIFICATION_DATASETS:
      seed = initial_seed
      for unused_seed_inc in range(num_seed):
        seed += 1
        dataset_info = datasets.DATASETS_SETTINGS[dataset_name]
        prior_knowledge = base.PriorKnowledge(
            input_dim=dataset_info.input_dim,
            num_train=dataset_info.num_train,
            num_classes=dataset_info.num_classes,
            tau=tau,
            temperature=temperature,
        )

        configs.append(
            ProblemConfig(
                prior_knowledge=prior_knowledge,
                seed=seed,
                dataset_name=dataset_name,
            ))
  return {f'classification{SEPARATOR}{i}': v for i, v in enumerate(configs)}

def classification_num_data_sweep(
    num_seed: int = 5,
    initial_seed: int = 0,
    temperature: float = 0.01,
) -> Dict[str, ProblemConfig]:
  """Generate hyperparameter sweep for classification problems with different number of training data.

  Args:
    num_seed: number of seeds per configuration of other hyperparameters.
    initial_seed: initial value of the seed.
    temperature: temperature to be used in prior_knowledge. It does not affect
      the data.

  Returns:
    Mapping problem_id: gp_settings (for use in gp_load).
  """
  configs = []
  for tau in [1, 10]:
    for dataset_name in datasets.CLASSIFICATION_DATASETS:
      seed = initial_seed
      for num_train in [1, 10, 100, 1000, 10_000, 100_000]:
        for unused_seed_inc in range(num_seed):
          seed += 1
          dataset_info = datasets.DATASETS_SETTINGS[dataset_name]
          # Update num_train of the dataset
          num_train = min(num_train, dataset_info.num_train)

          prior_knowledge = base.PriorKnowledge(
              input_dim=dataset_info.input_dim,
              num_train=num_train,
              num_classes=dataset_info.num_classes,
              tau=tau,
              temperature=temperature,
          )

          configs.append(
              ProblemConfig(
                  prior_knowledge=prior_knowledge,
                  seed=seed,
                  dataset_name=dataset_name,
              ))
  return {
      f'classification_variant_data{SEPARATOR}{i}': v
      for i, v in enumerate(configs)
  }

def regression_sweep(num_seed: int = 5,
                     initial_seed: int = 0,
                     noise_std: float = 1.) -> Dict[str, ProblemConfig]:
  """Generate hyperparameter sweep for regression problems.

  Args:
    num_seed: number of seeds per configuration of other hyperparameters.
    initial_seed: initial value of the seed.
    noise_std: noise_std to be used in prior_knowledge. It does not affect
      the data.

  Returns:
    Mapping problem_id: gp_settings (for use in gp_load).
  """
  configs = []
  for dataset_name in datasets.REGRESSION_DATASETS:
    seed = initial_seed
    for unused_seed_inc in range(num_seed):
      seed += 1
      dataset_info = datasets.DATASETS_SETTINGS[dataset_name]
      prior_knowledge = base.PriorKnowledge(
          input_dim=dataset_info.input_dim,
          num_train=dataset_info.num_train,
          num_classes=dataset_info.num_classes,
          tau=1,
          noise_std=noise_std,
      )

      configs.append(
          ProblemConfig(
              prior_knowledge=prior_knowledge,
              seed=seed,
              dataset_name=dataset_name,
          ))
  return {f'regression{SEPARATOR}{i}': v for i, v in enumerate(configs)}

SETTINGS = {**regression_sweep(),
            **classification_sweep(),
            **classification_num_data_sweep(),}

REGRESSION = tuple(regression_sweep().keys())
CLASSIFICATION = tuple(classification_sweep().keys())
CLASSIFICATION_NUM_DATA = tuple(classification_num_data_sweep().keys())


from keras.preprocessing.sequence import pad_sequences
from nltk import word_tokenize
import numpy as np
import keras
import json
import os

# Parameters of the model
MAX_SENTENCE_LENGTH = 60
OUTPUT_DIMENSION = 5
DATASET_FILEPATH = "Datasets/ner_dataset.txt"
WORD_INDEX_FILEPATH = "word_index.json"
MODEL_FILEPATH = "ner_trained_model.h5"

CLASS_MAP = {
    0: 'PER',
    1: 'ORG',
    2: 'LOC',
    3: 'MISC',
    4: 'O'
}

def _get_dataset(file):
    """
    Read the dataset in `file` line-by-line and return a list of sentences.

    :param file: A String containing the path to the text document containing data.

    return: A list of sentences, where each sentence is a list of tuples:
            (<current word>, <part of speech tag>, <chunk tag>, <name entity tag>)
    """

    dataset = []

    with open(file, "r") as file:

        sentence = []

        # Read all lines one by one
        for line in file:
            line_items = line.split()

            # Ignore DOCSTART lines
            if 'DOCSTART' in line:
                continue

            # If empty row, add sentence to dataset and re-initialise it
            if len(line_items) == 0:

                if len(sentence) > 0:
                    dataset.append(sentence)
                sentence = []
                continue

            # If row is not empty add word to sentence
            if len(line_items) > 0:
                sentence.append(tuple(line_items))

    return dataset

def _preprocess_outputs(dataset):
    """
    Creates a targets matrix of shape (NUM_SENTENCES, MAX_SENTENCE_LENGTH, OUTPUT_DIMENSION)

    :param dataset: A list of tuples (<current word>, <part of speech tag>, <chunk tag>, <name entity tag>)
    containing the whole data in our dataset

    returns: A numpy 3D matrix for the model's `targets`
    """

    NUM_SENTENCES = len(dataset)

    # Initialise targets to zeros
    targets = np.zeros((NUM_SENTENCES, MAX_SENTENCE_LENGTH, OUTPUT_DIMENSION))

    # Loop over each sentence in the dataset, and each word in sentences
    for j, sentence in enumerate(dataset):
        for i, word in enumerate(sentence):

            # Only add the first MAX_SENTENCE_LENGTH words
            if i < MAX_SENTENCE_LENGTH:

                # Assign a target output for each word in the sentence
                label = word[3]
                if 'PER' in label:
                    targets[j, i, :] = [1, 0, 0, 0, 0]
                elif 'ORG' in label:
                    targets[j, i, :] = [0, 1, 0, 0, 0]
                elif 'LOC' in label:
                    targets[j, i, :] = [0, 0, 1, 0, 0]
                elif 'MISC' in label:
                    targets[j, i, :] = [0, 0, 0, 1, 0]
                elif label == 'O':
                    targets[j, i, :] = [0, 0, 0, 0, 1]
                else:
                    targets[j, i, :] = [0, 0, 0, 0, 0]

    return targets

def _preprocess_inputs(dataset_, word_index):
    """
    Loops through all words in all sentences and:
    (1) converts them to indices
    (2) adds padding to a length of MAX_SENTENCE_LENGTH

    :param dataset_: A list of tuples (<current word>, <part of speech tag>, <chunk tag>, <name entity tag>)
    containing the whole data in our dataset
    :param word_index: A dictionary with (<word>, <integer number>) (key, value) pairs

    return: A numpy matrix for the inputs with shape (NUM_SENTENCES, MAX_SENTENCE_LENGTH)
    """

    # initialise a list for the input data
    inputs = []

    # Loop through sentences
    for i, sentence in enumerate(dataset_):
        sentence = [word[0] for word in sentence]
        input_ = []

        # Loop through words in sentence
        for word in sentence:
            word = word.lower()
            word_num = word_index.get(word)
            if word_num is None:
                word_num = word_index.get('unk')
            input_.append(word_num)
        inputs.append(input_)

    # Pad sentences to some maximum length. Padding and truncation is done at the back of the sentence
    inputs = pad_sequences(inputs, maxlen=MAX_SENTENCE_LENGTH, padding='post', truncating='post')

    return inputs

def _load_trained_model(filepath='ner_trained_model.h5'):
    if not os.path.isfile(filepath):
        print("'ner_trained_model.h5' file is missing! Please put the file in the project's main folder.")
        return None
    return keras.models.load_model(filepath)

def _load_word_index(filepath='word_index.json'):
    if not os.path.isfile(filepath):
        print("'word_index.json' file is missing! Please put the file in the project's main folder.")
        return None

    with open(filepath, 'r') as file:
        word_index = json.load(file)
    return word_index

def _calculate_metrics_sentence(target_sentence, pred_sentence):

    # Filter out padding - this is where the target word is [0, 0, 0, 0, 0]
    for index, word in enumerate(target_sentence):
        if np.sum(word == np.array([0, 0, 0, 0, 0])) == OUTPUT_DIMENSION:
            break
    target_sentence = target_sentence[:index]
    pred_sentence = pred_sentence[:index]

    # Initialise results dictionary
    results = {}
    classes = list(range(target_sentence.shape[1]))
    for class_ in classes:
        results[class_] = {'tp': 0, 'fp': 0, 'fn': 0}

    # Get indices of target and predicted name entity class
    target_indices = np.argmax(target_sentence, axis=1)
    pred_indices = np.argmax(pred_sentence, axis=1)

    # Go through each target and prediction and sum up metrics
    for targ, pred in zip(target_indices,pred_indices):
        if targ == pred:
            results[targ]['tp'] += 1
            continue
        else:
            results[targ]['fn'] += 1
            results[pred]['fp'] += 1

    return results

def _calculate_metrics_total(model, test_data):
    test_x = test_data[0]
    test_y = test_data[1]

    preds = model.predict(test_x)

    # Initialise results dictionary
    results = {}
    classes = list(range(test_y[0].shape[1]))
    for class_ in classes:
        results[class_] = {'tp': 0, 'fp': 0, 'fn': 0}

    # Calculate metrics for each sentence one by one and add up results
    for target_sentence, pred_sentence in zip(test_y, preds):
        res_sentence = _calculate_metrics_sentence(target_sentence, pred_sentence)

        # Sum up results
        for class_ in classes:
            results[class_]['tp'] += res_sentence[class_]['tp']
            results[class_]['fp'] += res_sentence[class_]['fp']
            results[class_]['fn'] += res_sentence[class_]['fn']

    return results

def predict_on_test_set(filepath):

    # Calculate and print metrics
    results = _calculate_metrics_total(model, (inputs, targets))
    for label in results.keys():
        precision = results[label]['tp'] / (results[label]['tp'] + results[label]['fp'])
        recall = results[label]['tp'] / (results[label]['tp'] + results[label]['fn'])
        f1_score = 2 / (1 / precision + 1 / recall)
        support = 0
        for target in targets:
            support += np.sum(np.argmax(target) == label)
        print("%4s: precision = %5.3f recall = %5.3f f1_score = %5.3f support = %6d" %
              (
                  CLASS_MAP[label],
                  precision,
                  recall,
                  f1_score,
                  support
              ))

def _sentence_to_vector(sentence):

    # Import our word dictionary and tokenize the sentence
    word_index = _load_word_index(filepath='word_index2.json')
    tokens = word_tokenize(sentence)

    # Map word tokens into a vector of numbers
    tokens_num = [word_index.get(token.lower(), 0) for token in tokens]
    vector = np.zeros(shape=(1, MAX_SENTENCE_LENGTH))
    vector[0, : len(tokens_num)] = tokens_num

    return vector, tokens

def predict_on_texts(texts):

    for sentence in texts:
        vector, clean_sentence = _sentence_to_vector(sentence)
        preds = model.predict(vector)
        preds = preds[:, :len(clean_sentence)]
        pred_output_classes = np.argmax(preds[0], axis=1)

        for word, pred in zip(clean_sentence, pred_output_classes):
            print(word, end=" ")
            print("[%s]" % CLASS_MAP[pred], end=" ")
        print()

"""
------------------------------------------------------------------
The following code will run with importing the script
------------------------------------------------------------------
"""

# Load trained model
model = _load_trained_model(MODEL_FILEPATH)

# Create a dataset from text file
dataset = _get_dataset(DATASET_FILEPATH)

# Load dictionaries for mapping words to numbers and vice versa
word_index = _load_word_index(filepath=WORD_INDEX_FILEPATH)

# Process output labels
targets = _preprocess_outputs(dataset)

# Process input data
inputs = _preprocess_inputs(dataset, word_index)

"""
------------------------------------------------------------------
"""

if __name__ == '__main__':

    """
    Test the script here...
    """

    pass
    # texts = ["Hello how $34 are you doing, have you been to miami before? How about tokyo!?",
    #          "My name is Mihajlo, what's your name?",
    #          "I work at Imperial College London, which is located in London, UK",
    #          "Inter will beat Juventus 3:0 on Friday"]

    # predict_on_texts(texts)
    # predict_on_test_set(filepath=DATASET_FILEPATH)


import numpy as np
from GCN_ETHZ.vocabulary import Vocabulary
import torch
from torch.autograd import Variable
import GCN_ETHZ.dataset as D
import GCN_ETHZ.utils as utils
import GCN_ETHZ.ntee as ntee
from GCN_ETHZ.gcn.model import GCN
from random import shuffle
import torch.optim as optim
from GCN_ETHZ.abstract_word_entity import load as load_model
from GCN_ETHZ.mulrel_ranker import MulRelRanker
from pprint import pprint
from itertools import count
import copy
import csv
import json
import time
from collections import Counter
# import ipdb
# import pdb
import random
import math

ModelClass = MulRelRanker
wiki_prefix = 'en.wikipedia.org/wiki/'
debugging = False
debugging_skip_train = False

class EDRanker:
    """
    ranking candidates
    """
    def __init__(self, config):
        print('--- create EDRanker model ---')

        config['entity_embeddings'] = config['entity_embeddings'] / \
                                      np.maximum(np.linalg.norm(config['entity_embeddings'],
                                                                axis=1, keepdims=True), 1e-12)
        config['entity_embeddings'][config['entity_voca'].unk_id] = 1e-10
        config['word_embeddings'] = config['word_embeddings'] / \
                                    np.maximum(np.linalg.norm(config['word_embeddings'],
                                                              axis=1, keepdims=True), 1e-12)
        config['word_embeddings'][config['word_voca'].unk_id] = 1e-10
        self.one_entity_once = config['one_entity_once']
        self.ent_inlinks = config['entity_inlinks']
        self.n_sample = config['n_sample']
        self.search_ment_size = config['search_ment_size']
        self.search_entity_size = config['search_entity_size']
        self.predict_epoches = config['predict_epoches']
        self.death_epoches = config['death_epoches']
        self.local_stop_epoches = config['local_stop_epoches']
        self.change_rate = config['change_rate']
        self.local_cur_stop_epoches = 0

        self.word_vocab = config['word_voca']
        self.ent_vocab = config['entity_voca']

        self.output_path = config['f1_csv_path']
        print('prerank model')
        self.prerank_model = ntee.NTEE(config)
        self.args = config['args']

        print('main model')
        if self.args.mode == 'eval' or self.args.mode == 'load_train':
            print('try loading model from', self.args.model_path)
            self.model = load_model(self.args.model_path, ModelClass)
        else:
            print('create new model')

            config['use_local_epoches'] = self.args.use_local_epoches
            config['oracle'] = False
            self.model = ModelClass(config)

        self.load_ent_desc(500, 3)

        self.prerank_model.cuda()
        self.emb_dims = self.model.emb_dims
        self.negsam_graph_cache = {}
        self.true_cands = {}
        # self.model = torch.nn.DataParallel(self.model)
        self.model.cuda()

    def load_ent_desc(self, max_desc_len, n_grams):
        ent_desc = json.load(open('../data/ent2desc.json', 'r'))
        self.ent_desc = [[self.word_vocab.get_id(Vocabulary.unk_token) for j in range(max_desc_len)] for i in range(self.ent_vocab.size())]
        self.desc_mask = [[0 for j in range(max_desc_len-n_grams+1)] for i in range(self.ent_vocab.size())]
        for ent in ent_desc:
            for i in range(min(len(ent_desc[ent]), max_desc_len)):
                self.ent_desc[self.ent_vocab.get_id(ent)][i] = self.word_vocab.get_id(ent_desc[ent][i])
                if (i>=n_grams-1):
                    # seems that only pick the last n-gram words?
                    self.desc_mask[self.ent_vocab.get_id(ent)][i-(n_grams-1)] = 1
        self.ent_desc = Variable(torch.LongTensor(self.ent_desc).cuda())
        self.desc_mask = Variable(torch.LongTensor(self.desc_mask).cuda())

    def get_data_items(self, dataset, predict=False, isTrain=False):
        data = []
        cand_source = 'candidates'

        for doc_name, content in dataset.items():
            items = []

            for m in content:
                try:
                    named_cands = [c[0] for c in m[cand_source]]
                    p_e_m = [min(1., max(1e-3, c[1])) for c in m[cand_source]]
                    etype = [c[2] for c in m[cand_source]]
                except:
                    named_cands = [c[0] for c in m['candidates']]
                    p_e_m = [min(1., max(1e-3, c[1])) for c in m['candidates']]
                    etype = [c[2] for c in m['candidates']]
                try:
                    # true_pos is the index of true answer in the list of candidates (named_cands)
                    true_pos = named_cands.index(m['gold'][0])
                    p = p_e_m[true_pos]
                except:
                    true_pos = -1

                # select # n_cands_before_rank candidates according to their priority
                named_cands = named_cands[:min(self.args.n_cands_before_rank, len(named_cands))]
                p_e_m = p_e_m[:min(self.args.n_cands_before_rank, len(p_e_m))]
                etype = etype[:min(self.args.n_cands_before_rank, len(etype))]
                # guarantee that the ground truth is in the top30 candidates
                ### why they put the last top-cand as the fake truePos?
                if true_pos >= len(named_cands):
                    if not predict:
                        true_pos = len(named_cands) - 1
                        p_e_m[-1] = p
                        named_cands[-1] = m['gold'][0]
                    else:
                        true_pos = -1

                if true_pos < 0 and not predict:
                    true_pos = 0

                # cands is the dict_index of candidates
                cands = [self.model.entity_voca.get_id(wiki_prefix + c) for c in named_cands]
                mask = [1.] * len(cands)

                if len(cands) == 0 and not predict:
                    continue
                # if cands is not enough, pad it
                elif len(cands) < self.args.n_cands_before_rank:
                    cands += [self.model.entity_voca.unk_id] * (self.args.n_cands_before_rank - len(cands))
                    etype += [[0, 0, 0, 1]] * (self.args.n_cands_before_rank - len(etype))
                    named_cands += [Vocabulary.unk_token] * (self.args.n_cands_before_rank - len(named_cands))
                    p_e_m += [1e-8] * (self.args.n_cands_before_rank - len(p_e_m))
                    mask += [0.] * (self.args.n_cands_before_rank - len(mask))

                lctx = m['context'][0].strip().split()
                lctx_ids = [self.prerank_model.word_voca.get_id(t) for t in lctx if utils.is_important_word(t)]
                lctx_ids = [tid for tid in lctx_ids if tid != self.prerank_model.word_voca.unk_id]
                lctx_ids = lctx_ids[max(0, len(lctx_ids) - self.args.ctx_window//2):]

                rctx = m['context'][1].strip().split()
                rctx_ids = [self.prerank_model.word_voca.get_id(t) for t in rctx if utils.is_important_word(t)]
                rctx_ids = [tid for tid in rctx_ids if tid != self.prerank_model.word_voca.unk_id]
                rctx_ids = rctx_ids[:min(len(rctx_ids), self.args.ctx_window//2)]

                # seem to be one word at most time?
                # it's 'mention' with conll_m but not 'mentions' in conll_doc
                ment = m['mention'].strip().split()
                ment_ids = [self.prerank_model.word_voca.get_id(t) for t in ment if utils.is_important_word(t)]
                ment_ids = [tid for tid in ment_ids if tid != self.prerank_model.word_voca.unk_id]

                m['sent'] = ' '.join(lctx + rctx)
                mtype = m['mtype']
                items.append({'context': (lctx_ids, rctx_ids),
                              'ment_ids': ment_ids,
                              'cands': cands,
                              'named_cands': named_cands,
                              'p_e_m': p_e_m,
                              'mask': mask,
                              'true_pos': true_pos,
                              'mtype': mtype,
                              'etype': etype,
                              'doc_name': doc_name,
                              'raw': m,
                              'prev_dist': m['prev_dist']
                              })

            # one doc one time here
            if len(items) > 0:
                # note: this shouldn't affect the order of prediction because we use doc_name to add predicted entities,
                # and we don't shuffle the data for prediction

                # # ----old implementation----
                # # each doc is regarded as one batch
                # # data.append(items)
                #     if isTrain:
                #         for k in range(0, len(items), self.seq_len // 2):
                #             data.append(items[max(0, k - self.seq_len//2) : min(len(items), k + self.seq_len//2)])
                #     else:
                #         if self.one_entity_once:
                #             for k in range(0, len(items)):
                #                 data.append(items[max(0, k-self.seq_len+1): k+1])
                #         else:
                #             for k in range(0, len(items), self.seq_len):
                #                 data.append(items[k:min(len(items), k + self.seq_len)])

                # in order to use GCN, we try to process one whole doc for one time
                ### New ###
                data.append(items)
                ###     ###

        # every element in 'data' is a list of items
        return self.prerank(data, predict)

    # executed between the get_data_item and training
    def prerank(self, dataset, predict=False):
        new_dataset = []
        has_gold = 0
        total = 0

        for content in dataset:
            # content is a list of items
            items = []

            if self.args.keep_ctx_ent > 0:
                # rank the candidates by ntee scores
                lctx_ids = [m['context'][0][max(len(m['context'][0]) - self.args.prerank_ctx_window // 2, 0):]
                            for m in content]
                rctx_ids = [m['context'][1][:min(len(m['context'][1]), self.args.prerank_ctx_window // 2)]
                            for m in content]
                ment_ids = [[] for m in content]

                token_ids = [l + m + r if len(l) + len(r) > 0 else [self.prerank_model.word_voca.unk_id]
                             for l, m, r in zip(lctx_ids, ment_ids, rctx_ids)]

                entity_ids = [m['cands'] for m in content]
                entity_ids = Variable(torch.LongTensor(entity_ids).cuda())

                entity_mask = [m['mask'] for m in content]
                entity_mask = Variable(torch.FloatTensor(entity_mask).cuda())

                # token_ids' element is a part of the document which around the target mention
                token_ids, token_offsets = utils.flatten_list_of_lists(token_ids)
                token_ids = Variable(torch.LongTensor(token_ids).cuda())
                token_offsets = Variable(torch.LongTensor(token_offsets).cuda())

                # log_probs seem to be n_ments * n_cands ?
                log_probs = self.prerank_model.forward(token_ids, token_offsets, entity_ids, use_sum=True)
                log_probs = (log_probs * entity_mask).add_((entity_mask - 1).mul_(1e10))
                # top_pos seem to be n_ments * k ?
                _, top_pos = torch.topk(log_probs, dim=1, k=self.args.keep_ctx_ent)
                top_pos = top_pos.data.cpu().numpy()
            else:
                top_pos = [[]] * len(content)

            # select candidats: mix between keep_ctx_ent best candidates (ntee scores) with
            # keep_p_e_m best candidates (p_e_m scores)
            for i, m in enumerate(content):
                sm = {'cands': [],
                      'named_cands': [],
                      'p_e_m': [],
                      'mask': [],
                      'etype': [],
                      'true_pos': -1}
                m['selected_cands'] = sm

                selected = set(top_pos[i])
                idx = 0
                # insert remained candidate until fit the requirement
                while len(selected) < self.args.keep_ctx_ent + self.args.keep_p_e_m:
                    if idx not in selected:
                        selected.add(idx)
                    idx += 1

                selected = sorted(list(selected))
                for idx in selected:
                    if idx>len(m['cands'])-1:
                        continue
                    sm['cands'].append(m['cands'][idx])
                    sm['named_cands'].append(m['named_cands'][idx])
                    sm['p_e_m'].append(m['p_e_m'][idx])
                    sm['mask'].append(m['mask'][idx])
                    sm['etype'].append(m['etype'][idx])
                    # record the answer's position in sm
                    if idx == m['true_pos']:
                        sm['true_pos'] = len(sm['cands']) - 1

                if not predict:
                    if sm['true_pos'] == -1:
                        continue

                items.append(m)
                if sm['true_pos'] >= 0:
                    has_gold += 1
                total += 1

                if predict:
                    # only for oracle model, not used for eval
                    # h? why 'predict' means oracle model but not eval?
                    if sm['true_pos'] == -1:
                        sm['true_pos'] = 0  # a fake gold, happens only 2%, but avoid the non-gold

            if len(items) > 0:
                new_dataset.append(items)

        print('recall', has_gold / total)
        return new_dataset

    def gold_e_graph_build(self, dataset):
        cand_to_idxs = []
        idx_to_cands = []
        gold_e_adjs = []
        for dc, batch in enumerate(dataset):
            sele_cand = [m['selected_cands']['cands'] for m in batch]
            true_pos = [m['selected_cands']['true_pos'] for m in batch]
            true_cands = [sele_cand[idx][true_pos[idx]] if true_pos[idx] > -1 else sele_cand[idx][0] for idx in range(len(true_pos))]
            self.true_cands[dc] = true_cands
            cand_to_idx, idx_to_cand, e_adj = self.e_graph_build(true_cands)
            cand_to_idxs.append(cand_to_idx)
            idx_to_cands.append(idx_to_cand)
            gold_e_adjs.append(e_adj)
        return cand_to_idxs, idx_to_cands, gold_e_adjs

    def e_graph_build(self, cand_ids, hops=2, max_nm=False):
        if max_nm == False:
            max_nm = self.args.max_nm
        cand_to_idx = {}
        idx_to_cand = copy.deepcopy(cand_ids)
        node_counter = Counter()
        max_cand = max(self.ent_inlinks.keys())
        cnt = 0
        n_ments = len(cand_ids)
        max_n_node = int(max_nm / n_ments)
        if hops == 2:
            for c in cand_ids:
                if c not in cand_to_idx:
                    cand_to_idx[c] = []
                    if c in self.ent_inlinks:
                        neighbor = self.ent_inlinks[c]
                        node_counter.update(neighbor)
                cand_to_idx[c].append(cnt)
                cnt+=1
#                 if c in self.ent_inlinks:
#                     neighbor = self.ent_inlinks[c]
#                     node_counter.update(neighbor)
            for n in list(node_counter.elements()):
                if node_counter[n] > 1 and (n not in cand_to_idx) and n < self.model.entity_voca.size():
                    if n not in cand_to_idx:
                        cand_to_idx[n] = []
                    cand_to_idx[n].append(cnt)
                    idx_to_cand.append(n)
                    cnt+=1
                    if cnt >= max_n_node:
                        break
            n = len(idx_to_cand)
            e_adj = np.zeros((n, n))
            for cand, idx in cand_to_idx.items():
                n_idx = len(idx)
                slic = e_adj[idx]
                slic[:,idx] = np.ones((n_idx,n_idx))-np.eye(n_idx)
                e_adj[idx] = slic
                if cand not in self.ent_inlinks:
                    # print(cand, "not in ent_inlinks")
                    continue
                neighbor = [nid for a in self.ent_inlinks[cand] if a in cand_to_idx for nid in cand_to_idx[a]]
                slic = e_adj[idx]
                slic[:,neighbor]=1
                e_adj[idx] = slic
                slic = e_adj[neighbor]
                slic[:,idx]=1
                e_adj[neighbor] = slic
            return cand_to_idx, idx_to_cand, e_adj
        elif hops == 1:
            for c in cand_ids:
                cand_to_idx[c] = len(cand_to_idx)
            n = len(idx_to_cand)
            e_adj = np.zeros((n, n))
            for cand, idx in cand_to_idx.items():
                if cand not in self.ent_inlinks:
                    continue
                neighbor = [cand_to_idx[a] for a in self.ent_inlinks[cand] if a in cand_to_idx]
                e_adj[idx, neighbor] = 1
                e_adj[neighbor, idx] = 1
            return cand_to_idx, idx_to_cand, e_adj

    def ment_neg_sample(self, n_sample, entity_ids, true_pos, entity_mask):
        # entity_ids is n_ment * n_cand(k) indexes
        n_ment, n_cand = entity_ids.size()

        copy_entity_mask = entity_mask.index_put((torch.LongTensor(np.arange(n_ment)).cuda(), true_pos.long()), torch.zeros(n_ment).cuda())
        only_one_cand = (copy_entity_mask.sum(dim=1)==0).float().cuda()
        copy_entity_mask = entity_mask.index_put((torch.LongTensor(np.arange(n_ment)).cuda(), true_pos.long()), only_one_cand)
        sample_idx = torch.multinomial(copy_entity_mask, n_sample, replacement=True)

        i_matrix = torch.LongTensor(list(range(n_ment))).cuda().unsqueeze(1).repeat(1, n_sample)
        only_one_cand = only_one_cand.unsqueeze(1).repeat(1, n_sample)
        copy_entity_mask.index_put_((torch.flatten(i_matrix), torch.flatten(sample_idx)), torch.flatten(only_one_cand))
        # sample_idx is n_ment * n_sample
        # copy_entity_mask: n_ment * n_cand
        return sample_idx, copy_entity_mask

#     def ment_neg_sample(self, n_sample, entity_ids, true_pos, entity_mask, change_rate):
#         # entity_ids is n_ment * n_cand(k) indexes
#         n_ment, n_cand = entity_ids.size()
#         ment_res_num = torch.FloatTensor([n_sample] * n_ment).cuda()
#         assert change_rate <= 1.0
#         change_one_time = max(int(change_rate * n_ment), 1)
#         copy_entity_mask = entity_mask.index_put((torch.LongTensor(np.arange(n_ment)).cuda(), true_pos.long()), torch.zeros(n_ment).cuda())
#         only_one_cand = (copy_entity_mask.sum(dim=1)==0).float().cuda()
#         copy_entity_mask = entity_mask.index_put((torch.LongTensor(np.arange(n_ment)).cuda(), true_pos.long()), only_one_cand)

#         neg_sample_idx = true_pos.unsqueeze(0).unsqueeze(0).repeat(n_ment, n_sample, 1)
#         copy_entity_ids = entity_ids.t().squeeze(0).repeat(n_ment, 1, 1)
#         # neg_sample_idx: n_ment * n_sample * n_ment
#         # copy_entity_ids: n_ment * n_cand * n_ment

#         while ment_res_num.sum() > 0:
#             sample_ment_idx = torch.multinomial(ment_res_num, min(change_one_time, (ment_res_num>0).sum()), replacement=False)
#             sample_mask = copy_entity_mask[sample_ment_idx]
#             sample_ent_idx = torch.multinomial(sample_mask, 1, replacement=False).squeeze(1)
# #             print(sample_ment_idx, sample_ent_idx)
#             neg_sample_idx[:,:,sample_ment_idx] = neg_sample_idx[:,:,sample_ment_idx].index_put((sample_ment_idx, n_sample-ment_res_num[sample_ment_idx].long()), sample_ent_idx)
#             #print(neg_sample_idx[:,:,sample_ment_idx])
#             ment_res_num[sample_ment_idx] -= 1

#         return neg_sample_idx, torch.gather(copy_entity_ids, 1, neg_sample_idx)

    # Heuristic Order Learning Method - Based on Mention-Local Similarity or Mention-Topical Similarity
    def train(self, org_train_dataset, org_dev_datasets, config):
        print('extracting training data')
        org_train_dataset, train_mlist, train_madj = org_train_dataset
        train_dataset = self.get_data_items(org_train_dataset, predict=False, isTrain=True)

        doc_names = list(train_mlist.keys())
        shuffle_list = list(zip(train_dataset, doc_names))
        shuffle(shuffle_list)
        train_dataset[:], doc_names[:] = zip(*shuffle_list)
        train_dataset_adj = utils.data_m_graph_build(train_dataset)

        gold_cand_to_idxs, gold_idx_to_cand, gold_e_adjs = self.gold_e_graph_build(train_dataset)
        print('#train docs', len(train_dataset))
        self.init_lr = config['lr']
        dev_datasets = []
        for dname, data, mlist, madj in org_dev_datasets:
            dataitems = self.get_data_items(data, predict=True, isTrain=False)
            dev_datasets.append((dname, dataitems, mlist, utils.data_m_graph_build(dataitems)))
            print(dname, '#dev docs', len(dev_datasets[-1][1]))

        print('creating optimizer')
        optimizer_local = optim.Adam(self.model.local_parameter(), lr=config['lr'])
        optimizer_global = optim.Adam(self.model.global_parameter(), lr=config['lr'], weight_decay=0.02)

        for param_name, param in self.model.named_parameters():
            if param.requires_grad:
                print(param_name)

        best_f1 = -1
        not_better_count = 0
        is_counting = True
        eval_after_n_epochs = self.args.eval_after_n_epochs

        order_learning = False
        # order_learning_count = 0

        rl_acc_threshold = 0.7

        # optimize the parameters within the disambiguation module first
        # self.model.switch_order_learning(0)
        best_aida_A_rlts = []
        best_aida_A_f1 = 0.
        best_aida_B_rlts = []
        best_aida_B_f1 = 0.
        best_ave_rlts = []
        best_ave_f1 = 0.

        first_global_flag = True

        self.run_time = []
        for e in range(config['n_epochs']):

            if debugging_skip_train:
                for di, (dname, data, mlist, madj) in enumerate(dev_datasets):
                    if dname == 'aida-B':
                        self.rt_flag = True
                    else:
                        self.rt_flag = False
                    # pdb.set_trace()
                    predictions = self.predict(data, mlist, madj, isLocal=False)
                    #self.records[e][dname] = self.record
                    f1, precisioin = D.eval(org_dev_datasets[di][1], predictions)

                    print(dname, 'micro F1:', str(f1), ", precisioin:", str(precisioin), flush=True)
                break

            total_loss = 0
            cur_max_n_ment = 0
            start_time = time.time()
            isLocal = self.local_only_judge(e)

            for dc, batch in enumerate(train_dataset):  # each document is a minibatch
                self.model.train()
                # print("dc:",dc,"start")
                if first_global_flag:
                    dc_start_time = time.time()

                # convert data items to pytorch inputs
                token_ids = [m['context'][0] + m['context'][1]
                             if len(m['context'][0]) + len(m['context'][1]) > 0
                             else [self.model.word_voca.unk_id]
                             for m in batch]

                ment_ids = [m['ment_ids'] if len(m['ment_ids']) > 0
                            else [self.model.word_voca.unk_id]
                            for m in batch]

                entity_ids = Variable(torch.LongTensor([m['selected_cands']['cands'] for m in batch]).cuda())
                true_pos = Variable(torch.LongTensor([m['selected_cands']['true_pos'] for m in batch]).cuda())
                p_e_m = Variable(torch.FloatTensor([m['selected_cands']['p_e_m'] for m in batch]).cuda())
                entity_mask = Variable(torch.FloatTensor([m['selected_cands']['mask'] for m in batch]).cuda())

                # entity_ids is n_ment * n_cand(k) indexes
                # for every entity candidate, we have a description
                desc_ids = torch.index_select(self.ent_desc, 0, entity_ids.view(-1)).view(entity_ids.size(0), entity_ids.size(1), -1)
                desc_mask = torch.index_select(self.desc_mask, 0, entity_ids.view(-1)).view(entity_ids.size(0), entity_ids.size(1), -1)

                mtype = Variable(torch.FloatTensor([m['mtype'] for m in batch]).cuda())
                etype = Variable(torch.FloatTensor([m['selected_cands']['etype'] for m in batch]).cuda())

                token_ids, token_mask = utils.make_equal_len(token_ids, self.model.word_voca.unk_id)
                token_ids = Variable(torch.LongTensor(token_ids).cuda())
                token_mask = Variable(torch.FloatTensor(token_mask).cuda())

                ment_ids, ment_mask = utils.make_equal_len(ment_ids, self.model.word_voca.unk_id)
                ment_ids = Variable(torch.LongTensor(ment_ids).cuda())
                ment_mask = Variable(torch.FloatTensor(ment_mask).cuda())

                optimizer_local.zero_grad()
                optimizer_global.zero_grad()
                # optimizer.zero_grad()

                n_ment, n_cand = entity_ids.size()
                cur_max_n_ment = max(cur_max_n_ment, n_ment)
                # print("cur_max_n_ment:", cur_max_n_ment)
                e_cand_to_idxs = [[] for _ in range(n_ment) ]
                e_idx_to_cands = [[] for _ in range(n_ment) ]
                e_adjs = [[] for _ in range(n_ment) ]
                sele_cand = [m['selected_cands']['cands'] for m in batch]

#                 # sample_idx: n_ment * n_sample
#                 # copy_entity_mask: n_ment * n_cand
                sample_idx, copy_entity_mask = self.ment_neg_sample(self.n_sample, entity_ids, true_pos, entity_mask)

                # neg_sample_idx: n_ment * n_sample * n_ment
                # neg_sample_cands: n_ment * n_sample * n_ment
                # sample_idx: n_ment * n_sample
#                 neg_sample_idx, neg_sample_cands = self.ment_neg_sample(self.n_sample, entity_ids, true_pos, entity_mask, self.change_rate)
#                 sample_idx = torch.gather(neg_sample_idx, 2, torch.LongTensor(range(n_ment)).cuda().unsqueeze(1).unsqueeze(2).repeat(1, self.n_sample, 1)).squeeze(2)

                if isLocal == False:
                    ###  USE LOCAL AND GLOBAL
                    for i in range(n_ment):
                        for j in range(self.n_sample):
                            true_cands = copy.deepcopy(self.true_cands[dc])
                            true_cands[i] = sele_cand[i][sample_idx[i][j]]
                            cand_to_idx, idx_to_cand, e_adj = self.e_graph_build(true_cands)
#                             cand_to_idx, idx_to_cand, e_adj = self.e_graph_build(neg_sample_cands[i][j].tolist())
                            e_cand_to_idxs[i].append(cand_to_idx)
                            e_idx_to_cands[i].append(idx_to_cand)
                            e_adjs[i].append(e_adj)
                        e_cand_to_idxs[i].append(gold_cand_to_idxs[dc])
                        e_idx_to_cands[i].append(gold_idx_to_cand[dc])
                        e_adjs[i].append(gold_e_adjs[dc])

                    # e_adjs = torch.LongTensor(e_adjs)
                    # e_adjs: n_ment * n_sample * n_entity * n_entity
                    new_adjs, new_node_cands, new_node_mask = utils.e_graph_batch_padding(e_cand_to_idxs, e_idx_to_cands, e_adjs, n_ment, self.n_sample)
                    # new_adjs: n_ment * (n_sample+1) * n_node * n_node
                    # new_node_cands: n_ment * (n_sample+1) * n_node
                    # new_node_mask: n_ment * (n_sample+1) * n_node
                    nega_e = (new_adjs, new_node_cands, new_node_mask)

                    cur_doc_name = doc_names[dc]

                    if first_global_flag:
                        dc_end_time = time.time()
                        dc_n_ment = new_node_mask.size(0)
                        dc_n_node = new_node_mask.size(2)
                        print("dc", dc, "n_node:", dc_n_node, ",n_ment:", dc_n_ment, ",n1m:", dc_n_node*dc_n_ment, ",n2m:", dc_n_node*dc_n_node*dc_n_ment, ",preparing use time:", dc_end_time-dc_start_time, flush=True)

                    scores, dual_scores = self.model.forward(token_ids, token_mask, entity_ids, entity_mask, p_e_m, mtype, etype, ment_ids, ment_mask, desc_ids, desc_mask, train_mlist[cur_doc_name], train_dataset_adj[dc], nega_e, sample_idx, gold=true_pos.view(-1, 1), isTrain=True, isLocal=False)
#                     if first_global_flag:
#                         dc_end_time = time.time()
#                         print("dc", dc,"total use time:", dc_end_time-dc_start_time, ",avg time:", (dc_end_time-start_time)/(dc+1), flush=True)
                    if dc == 0:
                        # print("global+local:",scores)
                        end_time = time.time()
                        print("dc0 local:",dual_scores[0])
                        print("global:",dual_scores[1])

                    # loss = self.model.loss(dual_scores[1], torch.LongTensor([self.n_sample]*n_ment).cuda(), isLocal=isLocal)
                    loss = self.model.loss(scores, torch.LongTensor([self.n_sample]*n_ment).cuda(), isLocal=isLocal)
                    loss.backward()
#                     if config['model_path'] != 'model_param9':
#                         optimizer_local.step()
                    optimizer_global.step()

                else:
                    ### USE LOCAL ONLY
                    scores, _ = self.model.forward(token_ids, token_mask, entity_ids, entity_mask, p_e_m, mtype, etype, ment_ids, ment_mask, desc_ids, desc_mask, False, False, False, sample_idx, gold=true_pos.view(-1, 1), isTrain=True, isLocal=True)

                    if self.local_cur_stop_epoches == self.local_stop_epoches - 1 and dc == 0 and (e + 1) % eval_after_n_epochs == 0:
                        print("final local scores:", scores)

                    loss = self.model.loss(scores, true_pos, isLocal=isLocal)
                    loss.backward()
                    optimizer_local.step()

                torch.nn.utils.clip_grad_norm_([p for p in self.model.parameters() if p.requires_grad], max_norm=40, norm_type=2)
#                 for (name, p) in self.model.named_parameters():
#                     if name == "att_mat_diag":
#                         # print(name, p.grad.data)
#                         if torch.isnan(p.grad.data).any():
#                             print("att_mat_diag NAN detected")
#                             ipdb.set_trace()
                            # p.grad.data.zero_()
#                             elif torch.sum(p.grad.data)==0:
#                                 print("sum 0 detected")
#                                 ipdb.set_trace()

                self.model.regularize(max_norm=2)

                loss = loss.cpu().data.numpy()
                total_loss += loss

                torch.cuda.empty_cache()

            end_time = time.time()
            print('epoch', e, 'total loss', total_loss, total_loss / len(train_dataset), "use time:", end_time-start_time, flush=True)
            if isLocal == False:
                first_global_flag = False

            if (e + 1) % eval_after_n_epochs == 0 or debugging:
                start_time = time.time()
                dev_f1 = 0.
                test_f1 = 0.
                ave_f1 = 0.
                if rl_acc_threshold < 0.92:
                    rl_acc_threshold += 0.02
                temp_rlt = []
                #self.records[e] = dict()
                for di, (dname, data, mlist, madj) in enumerate(dev_datasets):
                    if dname == 'aida-B':
                        self.rt_flag = True
                    else:
                        self.rt_flag = False
                    # predictions = self.predict(data, config['isDynamic'], order_learning)
                    predictions = self.predict(data, mlist, madj, isLocal)
                    #self.records[e][dname] = self.record
                    f1, precision = D.eval(org_dev_datasets[di][1], predictions)

                    if isLocal:
                        print("using local now, waiting for global later")
                    else:
                        print("using global now")
                    print(dname, 'micro F1: ' + str(f1), ', precision: ' + str(precision), flush=True)

                    with open(self.output_path, 'a') as eval_csv_f1:
                        eval_f1_csv_writer = csv.writer(eval_csv_f1)
                        eval_f1_csv_writer.writerow([dname, e, 0, f1])
                    temp_rlt.append([dname, f1])
                    if dname == 'aida-A':
                        dev_f1 = f1
                    if dname == 'aida-B':
                        test_f1 = f1
                    ave_f1 += f1
                if dev_f1>best_aida_A_f1:
                    best_aida_A_f1 = dev_f1
                    best_aida_A_rlts = copy.deepcopy(temp_rlt)
                if test_f1>best_aida_B_f1:
                    best_aida_B_f1 = test_f1
                    best_aida_B_rlts = copy.deepcopy(temp_rlt)
                if ave_f1 > best_ave_f1:
                    best_ave_f1 = ave_f1
                    best_ave_rlts = copy.deepcopy(temp_rlt)

                # if not config['isDynamic']:
                #     self.record_runtime('DCA')
                # else:
                if isLocal:
                    self.record_runtime('local')

                    if dev_f1 < best_f1:
                        self.local_cur_stop_epoches += 1
                    else:
                        self.local_cur_stop_epoches = 0
                        best_f1 = dev_f1
                        print('save model to', self.args.model_path)
                        self.model.save(self.args.model_path)
                else:
                    self.record_runtime('GCN')
                    print('save model to', self.args.model_path+'_'+str(e))
                    self.model.save(self.args.model_path+'_'+str(e))

                self.model.print_weight_norm()
                end_time = time.time()
                print("eval use time:", end_time-start_time)

        print('best_aida_A_rlts', best_aida_A_rlts)
        print('best_aida_B_rlts', best_aida_B_rlts)
        print('best_ave_rlts', best_ave_rlts)

    def record_runtime(self, method):
        self.run_time.sort(key=lambda x:x[0])
        pre_cands = 0
        count = 0
        total = 0.
        rt = dict()
        for cands, ti in self.run_time:
            if not cands==pre_cands:
                if pre_cands > 0:
                    rt[pre_cands] = total / count
                total = ti
                count = 1
                pre_cands = cands
            else:
                count += 1
                total += ti
        if count >0:
            rt[pre_cands] = total / count
        with open('runtime_%s.csv'%method, 'w') as runtime_csv:
            runtime_csv_writer = csv.writer(runtime_csv)
            for cands, ti in rt.items():
                runtime_csv_writer.writerow([cands, ti])
            runtime_csv.close()

    def predict(self, data, mlist, madj, isLocal=False):
        search_ment_size = self.search_ment_size
        search_entity_size = self.search_entity_size
        predict_epoches = self.predict_epoches
        death_epoches = self.death_epoches
        predictions = {items[0]['doc_name']: [] for items in data}
        self.model.eval()
        #self.record = []
        total_changed = 0
        total_ments = 0
        for dc, batch in enumerate(data):  # each document is a minibatch, is a list of mentions
            self.model.doc_predict_restore = False
            start_time = time.time()
            token_ids = [m['context'][0] + m['context'][1]
                         if len(m['context'][0]) + len(m['context'][1]) > 0
                         else [self.model.word_voca.unk_id]
                         for m in batch]

            ment_ids = [m['ment_ids'] if len(m['ment_ids']) > 0
                        else [self.model.word_voca.unk_id]
                        for m in batch]

            total_candidates = sum([len(m['selected_cands']['cands']) for m in batch])

            entity_ids = Variable(torch.LongTensor([m['selected_cands']['cands'] for m in batch]).cuda())
            p_e_m = Variable(torch.FloatTensor([m['selected_cands']['p_e_m'] for m in batch]).cuda())
            entity_mask = Variable(torch.FloatTensor([m['selected_cands']['mask'] for m in batch]).cuda())
            true_pos = Variable(torch.LongTensor([m['selected_cands']['true_pos'] for m in batch]).cuda())

            token_ids, token_mask = utils.make_equal_len(token_ids, self.model.word_voca.unk_id)

            token_ids = Variable(torch.LongTensor(token_ids).cuda())
            token_mask = Variable(torch.FloatTensor(token_mask).cuda())

            desc_ids = torch.index_select(self.ent_desc, 0, entity_ids.view(-1)).view(entity_ids.size(0),entity_ids.size(1), -1)
            desc_mask = torch.index_select(self.desc_mask, 0, entity_ids.view(-1)).view(entity_ids.size(0),entity_ids.size(1), -1)
            ment_ids, ment_mask = utils.make_equal_len(ment_ids, self.model.word_voca.unk_id)
            ment_ids = Variable(torch.LongTensor(ment_ids).cuda())
            ment_mask = Variable(torch.FloatTensor(ment_mask).cuda())

            mtype = Variable(torch.FloatTensor([m['mtype'] for m in batch]).cuda())
            etype = Variable(torch.FloatTensor([m['selected_cands']['etype'] for m in batch]).cuda())

            n_ments, n_cands = entity_ids.size()
            # the val in cur_cand_idxs should be in 0~n_cand
            maybe_no_cand = (entity_mask.sum(dim=1)==0)
            maybe_no_cand_idx = torch.arange(n_ments)[maybe_no_cand]
            if maybe_no_cand.any() and not isLocal:
#                 print("sum of entity_mask 0 detected")
                entity_mask[maybe_no_cand_idx,0] = 1
#                 print("mids:", maybe_no_cand_idx)
#                 print("eids:", entity_ids[maybe_no_cand])

            scores, _ = self.model.forward(token_ids, token_mask, entity_ids, entity_mask, p_e_m, mtype, etype, ment_ids, ment_mask, desc_ids, desc_mask, False, False, False, False, gold=False, isTrain=False, isLocal=True)
            final_pde = 0
            if isLocal:
                scores = scores.cpu().data.numpy()

                pred_ids = np.argmax(scores, axis=1)

            else:
                cur_cand_idxs = torch.argmax(scores, dim=1)
                local_cur_cands = torch.gather(entity_ids, 1, cur_cand_idxs.unsqueeze(1))

                death_cnt = 0
                print_flag = False # for debug
                for pde in range(predict_epoches):
                    cur_cands = torch.gather(entity_ids, 1, cur_cand_idxs.unsqueeze(1))
                    list_cands = cur_cands.squeeze(1).cpu().numpy().tolist()
                    # cur_cands: n_ment * 1
                    cand_to_idx, idx_to_cand, e_adj = self.e_graph_build(list_cands)
                    idx_to_cand = torch.LongTensor(idx_to_cand).cuda()
                    the_mask = torch.ones(idx_to_cand.size(0)).cuda()
                    nega_e = (torch.LongTensor(e_adj).cuda(), idx_to_cand, the_mask)
                    cur_scores, _ = self.model.forward(token_ids, token_mask, entity_ids, entity_mask, p_e_m, mtype, etype, ment_ids, ment_mask, desc_ids, desc_mask, None, madj[dc], nega_e, cur_cand_idxs, gold=None, isTrain=False, chosen_ment=False, isLocal=False)
                    assert cur_scores.size(0) == n_ments

                    if print_flag:
                        print_flag = False
                        print("changed new_scores:", cur_scores)

                    # small_scores, small_idxs = torch.topk(cur_scores.squeeze(1), min(search_ment_size, n_ments), largest=False, sorted=True)
                    small_idxs = torch.multinomial(torch.ones(n_ments).cuda(), min(search_ment_size, n_ments))
                    small_scores = cur_scores.squeeze(1)[small_idxs]

                    random_new_cand_idx = torch.cat([cur_cand_idxs[small_idxs].unsqueeze(1), torch.multinomial(entity_mask[small_idxs], min(n_cands, search_entity_size)-1, replacement=True)], dim=1)
                    # small_scores: search_ment_size
                    # random_new_cand_idx: search_ment_size * search_entity_size, value in 0~n_cands
                    e_cand_to_idxs = [[] for _ in range(min(search_ment_size, n_ments)) ]
                    e_idx_to_cands = [[] for _ in range(min(search_ment_size, n_ments)) ]
                    e_adjs = [[] for _ in range(min(search_ment_size, n_ments)) ]

                    for i in range(min(search_ment_size, n_ments)):
                        for j in range(min(n_cands, search_entity_size)):
                            true_cands = copy.deepcopy(list_cands)
                            which_ment = small_idxs[i]
                            true_cands[which_ment] = int(entity_ids[which_ment][random_new_cand_idx[i][j]])
                            cand_to_idx, idx_to_cand, e_adj = self.e_graph_build(true_cands)
                            e_cand_to_idxs[i].append(cand_to_idx)
                            e_idx_to_cands[i].append(idx_to_cand)
                            e_adjs[i].append(e_adj)
                    new_adjs, new_node_cands, new_node_mask = utils.e_graph_batch_padding(e_cand_to_idxs, e_idx_to_cands, e_adjs, min(search_ment_size, n_ments), min(n_cands, search_entity_size)-1)
                    nega_e = (new_adjs, new_node_cands, new_node_mask)

                    new_scores, _ = self.model.forward(token_ids, token_mask, entity_ids, entity_mask, p_e_m, mtype, etype, ment_ids, ment_mask, desc_ids, desc_mask, None, madj[dc], nega_e, random_new_cand_idx, gold=None, isTrain=False, chosen_ment=small_idxs, isLocal=False)
                    # new_scores: search_ment_size * search_entity_size

                    big_score_in_new, big_idxs_in_new = torch.max(new_scores, dim=1)
                    # big_idxs = random_new_cand_idx[big_idxs_in_new]
                    big_idxs = torch.gather(random_new_cand_idx, 1, big_idxs_in_new.unsqueeze(1)).squeeze(1)
                    final_pde = pde
                    if big_score_in_new.sum() - new_scores[:,0].sum() < 1e-7:
                        death_cnt = death_cnt + 1
                        if death_cnt > death_epoches:
                            if random.random()<0.01:
                                print("DEATH")
                                print("pde:", pde, "old_scores:", cur_scores)
                                print("new_scores:", new_scores)
                                print("------")
                            break
                    else:
#                         print("pde:", pde, "small_idxs:", small_idxs, "cur_cand_idxs[small_idxs]", cur_cand_idxs[small_idxs])
#                         print_flag = True
#                         print("before change old_scores:", cur_scores)
#                         print("big_idxs:", big_idxs)
                        cur_cand_idxs[small_idxs] = big_idxs
                        death_cnt = 0

                pred_ids = cur_cand_idxs
                global_cur_cands = torch.gather(entity_ids, 1, cur_cand_idxs.unsqueeze(1))
                total_ments += n_ments
                total_changed += int((global_cur_cands!=local_cur_cands).sum())

            end_time = time.time()
            #if isLocal == False and random.random()<0.1:
            if isLocal == False and final_pde > death_epoches:
                print("dc:", dc, "pred_time:", end_time-start_time, ",pde:", final_pde, ",n_ments:", n_ments)
            if isLocal == False and dc == 0:
                print("global_beta", self.model.global_beta, cur_scores.squeeze(1))

            if self.rt_flag:
                self.run_time.append([total_candidates, end_time-start_time])

            pred_entities = [m['selected_cands']['named_cands'][i] if m['selected_cands']['mask'][i] == 1
                             else (m['selected_cands']['named_cands'][0] if m['selected_cands']['mask'][0] == 1 else 'NIL')
                             for (i, m) in zip(pred_ids, batch)]

            doc_names = [m['doc_name'] for m in batch]
            self.added_words = []
            self.added_ents = []
#             if self.seq_len>0 and self.one_entity_once:
#                 predictions[doc_names[-1]].append({'pred': (pred_entities[-1], 0.)})
#             else:
            # for ids in self.model.added_words:
            #     self.added_words.append([self.word_vocab.id2word[idx] for idx in ids])
            # for ids in self.model.added_ents:
            #     self.added_ents.append([self.ent_vocab.id2word[idx] for idx in ids])
            for dname, entity in zip(doc_names, pred_entities):
                predictions[dname].append({'pred': (entity, 0.)})
            #self.record.append(dict({'added_words':self.added_words, 'added_ents':self.added_ents}))

        if isLocal==False:
            print("total changed:",total_changed,"total ments:",total_ments,"ratio:",total_changed/total_ments)
        return predictions

    def local_only_judge(self, cur_epoch):
        if self.local_cur_stop_epoches >= self.local_stop_epoches:
            return False
        if self.args.use_local_epoches < 0:
            return True
        if cur_epoch >= self.args.use_local_epoches:
            return False
        return True


from typing import List, Tuple, Dict

# Define types
Point = Tuple[int, int]
Points = List[Point]
Step = Tuple[str, int]
Path = List[Step]

def parse_input(text: str) -> Path:
    steps = text.split(",")
    return [(s[0], int(s[1:])) for s in steps]

def take_step(loc: Point, step: Step) -> Points:
    """
    create a list of points traversed in
    taking a step, in some direction
    """
    direction, length = step
    if direction is "U":
        return [(loc[0], loc[1] + i) for i in range(1, length + 1)]
    elif direction is "D":
        return [(loc[0], loc[1] - i) for i in range(1, length + 1)]
    elif direction is "R":
        return [(loc[0] + i, loc[1]) for i in range(1, length + 1)]
    else:
        return [(loc[0] - i, loc[1]) for i in range(1, length + 1)]

def points_from_path(loc: Point, path: Path) -> Points:
    """
    Turn a path into a list of points traversed on the path
    """
    points = [loc]
    for step in path:
        segment = take_step(loc, step)
        points.extend(segment)
        loc = segment[-1]

    return points

def intersections(points1: Points, points2: Points) -> Points:
    return list(set(points1[1:]) & set(points2[1:]))

def manhattan_distance_to_origin(pt: Point) -> int:
    return abs(pt[0]) + abs(pt[1])

def steps_taken(pt: Point, points1: Points, points2: Points) -> int:
    return points1.index(pt) + points2.index(pt)

if __name__ == "__main__":

    with open("day03.txt") as f:
        lines = f.read().strip().split("\n")

    paths = [parse_input(l) for l in lines]
    points = [points_from_path((0, 0), path) for path in paths]

    inters = intersections(*points)

    min_pt = min(inters, key=manhattan_distance_to_origin)
    print(min_pt, manhattan_distance_to_origin(min_pt))

    steps = lambda x: steps_taken(x, *points)
    min_steps = min(inters, key=steps)
    print(min_steps, steps_taken(min_steps, *points))


##4. Знайти максимальний елемент серед мінімальних елементів стовпців матриці.
from random import *
row = int(input("Enter the count of rows: "))
col = int(input("Enter the count of columns: "))
arr = [[0] * col for i in range(row)]
for i in range(row):
    for j in range(col):
        arr[i][j] = randint(0, 100)
        print("%4d "%arr[i][j], end = "")
    print()
print()

lst_min_elem = []
for j in range(col):
    temp_min = arr[i][j]
    for i in range(row):
        if temp_min > arr[i][j]:
            temp_min = arr[i][j]
    lst_min_elem.append(temp_min)
print("Min elems of every column:\n", lst_min_elem)
print("\nMax elem of min elems is : %d"%max(lst_min_elem))


# it is used for the purpose of adding security like encryption.
# make the ip address into some other format for encryption

def ip_address(address):
    # we declare one address variable
    new_address = ""
    # split the address using split function i
    split_address = address.split(" ")
    separator = "iswasthat"
    # join is used to merge separator and the split address
    new_address = separator.join(split_address)
    return new_address

ipaddress = ip_address("1 1 0 0 0 0 2 5 6")
print(ipaddress)