# Copyright 2016, Blur Studio
#
# This file is part of Simplex.
#
# Simplex is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Simplex is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.	See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Simplex.  If not, see <http://www.gnu.org/licenses/>.

# pylint: disable=invalid-name
from __future__ import division, absolute_import, print_function
import six
from six.moves import range, zip

try:
    import numpy as np
except ImportError:
    pass

INF = float("inf")
EPS = 1e-7

########################
# Mean Value Coords
########################

def _lerp(idx, corners, p):
    # If I'm here, I know that p lies on the line
    # between the corners at idx and idx+1.
    # Return the lerp value
    ip = (idx + 1) % len(corners)
    c1 = corners[idx]
    c2 = corners[ip]
    base = c1 - c2
    proj = np.dot(p - c2, base)
    b = proj / (base ** 2).sum()
    ret = np.zeros(len(corners))
    ret[idx] = b
    ret[ip] = 1.0 - b
    return ret

def mvc(corners, p, tol=EPS):
    """Get the Mean Value Coordinates of a point p in the polygon defined by corners

    Parameters
    ----------
    corners : [int, ...]
        A list of corner indices
    p : np.array
        The array of points
    tol : float
        A small tolerance value, defaulting to the global EPS

    Returns
    -------
    : list
        The normalized list of barycentric weights for this point in the polygon

    """
    spokes = corners - p
    spokeLens = (spokes * spokes).sum(axis=1)

    # Check if p is on top of a vertex
    for i, v in enumerate(spokeLens):
        if v < tol:
            bary = np.zeros(len(corners))
            bary[i] = 1.0
            return bary

    rspokes = np.roll(spokes, -1, axis=0)
    areas = np.cross(spokes, rspokes) * 0.5
    dots = (spokes * rspokes).sum(axis=1)

    for i, v in enumerate(areas):
        if v < tol and dots[i] < 0.0:
            return _lerp(i, corners, p)

    spokeLens = spokeLens ** 0.5
    rspokeLens = np.roll(spokeLens, -1)
    t = areas / (spokeLens * rspokeLens + dots)
    rawWeights = (np.roll(t, 1) + t) / spokeLens
    return rawWeights / sum(rawWeights)

def mmvc(rawFaces, points, samples, uvToFace, tol=EPS):
    """Multi-MeanValueCoords
        Get the MVC's of many points over many faces using numpy

    Parameters
    ----------
    rawFaces : [[int, ...], ...]
        A list of face index lists
    points : np.array
        A numpy array of uv points
    samples : np.array
        A numpy array of points to get the barycentric coords for
    uvToFace : {int: [int, ...], ...}
        A dictionary of UV indices to a list of Face Indices
    tol : float
        A small tolerance value, defaulting to the global EPS

    Returns
    -------
    : {fc: (wh, barys)}
        A dictionary of face counts to a tuple containing the indices of the face indices
        checked, and the barycentric coords

    """
    uvIdxs, faceIdxs = sorted(zip(*list(uvToFace.items())))
    uvIdxs = np.array(uvIdxs)
    faces = [rawFaces[fi] for fi in faceIdxs]
    faceLens = np.array([len(i) for i in faces])
    fcs = np.unique(faceLens)

    out = {}
    for fc in fcs:
        wh = np.where(faceLens == fc)[0]
        cIdxs = np.array([faces[i] for i in wh])
        barys = np.zeros(cIdxs.shape)
        qIdxs = uvIdxs[wh]
        out[fc] = (qIdxs, barys)  # edit barys in-place

        # cornerses is a [f, fc, 2] array
        # where f is num of faces, fc is verts per face, and 2 because uv is 2d
        # pts is [f, 2] array where f is num faces
        cornerses = points[cIdxs]
        pts = samples[qIdxs]

        # Get the "spokes" from the query point to the face corners
        # and get their squared lengths
        spokes = cornerses - pts[:, None, :]
        spokeLens2 = (spokes * spokes).sum(axis=-1)

        # Handle any samples that are directly on top of a corner
        # (where the spoke length is zero)
        zeros = np.any(spokeLens2 < tol, axis=1)
        onPoint = np.where(zeros)
        offPoint = np.where(~zeros)
        if onPoint[0].size:
            pIdxs = spokeLens2[onPoint].argmin(axis=1)
            barys[onPoint, pIdxs] = 1.0

        if not offPoint[0].size:
            continue

        # Ignore those points that have already been computed
        idxs = offPoint[0]
        spokes = spokes[offPoint]
        spokeLens2 = spokeLens2[offPoint]

        # Get the signed area of each triangle created by the spokes
        # and the dot product for the angle between each
        rspokes = np.roll(spokes, -1, axis=1)
        areas = np.cross(spokes, rspokes) * 0.5
        dots = (spokes * rspokes).sum(axis=-1)

        # Handle any samples that are directly on an edge between
        # two corners. (Where the triangle area is 0, and the
        # dot product is negative)
        aareas = abs(areas)
        zeros = aareas < tol
        ndots = dots < 0.0
        zn = zeros & ndots
        toLerp = np.any(zn, axis=-1)
        onEdge = np.where(toLerp)
        offEdge = np.where(~toLerp)
        if onEdge[0].size:
            subIdx = idxs[onEdge]
            pIdxs = zn[onEdge].argmax(axis=-1)
            xIdxs = (pIdxs + 1) % fc

            ppIdxs = cIdxs[subIdx, pIdxs]
            xpIdxs = cIdxs[subIdx, xIdxs]

            cp = pts[subIdx]
            pp = points[ppIdxs]
            xp = points[xpIdxs]

            bases = pp - xp
            diff = cp - xp
            proj = (bases * diff).sum(axis=-1)
            b = proj / (bases ** 2).sum(axis=1)
            mb = 1.0 - b
            barys[subIdx, pIdxs] = b
            barys[subIdx, xIdxs] = mb

        if not offEdge[0].size:
            continue

        idxs = idxs[offEdge]
        dots = dots[offEdge]
        areas = areas[offEdge]
        spokeLens2 = spokeLens2[offEdge]
        spokeLens = spokeLens2 ** 0.5
        rspokeLens = np.roll(spokeLens, -1, axis=1)

        t = areas / (spokeLens * rspokeLens + dots)
        rawWeights = (np.roll(t, -1, axis=1) + t) / spokeLens
        b = rawWeights / rawWeights.sum(axis=1)[..., None]
        barys[idxs] = b

    ret = {}
    for qIdxs, barys in six.itervalues(out):
        for qi, b in zip(qIdxs, barys):
            ret[qi] = (uvToFace[qi], b)

    return ret

########################
# Sweep algorithm
########################

def triArea(a, b, c):
    """Use the cross-ish-product to find the area of the triangle
        Depending on the winding, the area could be positive or negative

    Parameters
    ----------
    a : [float, float, float]
        The first point
    b : [float, float, float]
        The second point
    c : [float, float, float]
        The third point

    Returns
    -------
    : float
        The area of the triangle

    """
    return (a[0] * (c[1] - b[1]) + b[0] * (a[1] - c[1]) + c[0] * (b[1] - a[1])) / 2.0

def pointInTri(p, a, b, c, tol=EPS):
    """Check that the point is inside the triangle

    Parameters
    ----------
    p : [float, float, float]
        The point to check is inside the triangle
    a : [float, float, float]
        The first point of the triangle
    b : [float, float, float]
        The second point of the triangle
    c : [float, float, float]
        The third point of the triangle
    tol : float
        A small tolerance value, defaulting to the global EPS

    Returns
    -------
    : bool
        Whether the point is inside the given triangle

    """
    area = abs(triArea(a, b, c))
    chis = abs(triArea(p, b, c)) + abs(triArea(a, p, c)) + abs(triArea(a, b, p))
    return abs(area - chis) < tol

def sweep(qPoints, uvs, tris, pBar=None):
    """Get what triangle each query point is inside

    Runs a sweep-line algorithm to check for points in triangles.
    Imagine a uv layout, and a bunch of points on the layout.
    Now imagine a vertical line sweeping across the uv plane from the left
    to the right. The sorted u-values of certain properties then handled
    one-by-one.

    The algorithm keeps track of what triangles the vertical line is currently intersecting.

    * This is done by storing the min and max u-values for each triangle.
    * If the u-value encountered is the min for a tri, that tri is added to the list
    * If the u-value encountered is the max for a tri, that tri is removed from the list

    Then, each time a query point is encountered, it only has to check the intersection list

    Parameters
    ----------
    qPoints : np.array
        The points to get barycentric coordinates of
    uvs : np.array
        The UVs we're searching
    tris : np.array
        The Nx3 array of triangle indices
    pBar : QProgressDialog, optional
        An optional progress dialog

    Returns
    -------
    : {int: int, ...}
        A dictionary of out[pointIndex] = triIndex
    : set()
        A set containing the pointIndexes that weren't found in a triangle

    """
    tpts = uvs[tris]  # [tIdx, abc, xy]

    allmxs, allmns = tpts.max(axis=1), tpts.min(axis=1)
    ymxs, ymns = allmxs[:, 1], allmns[:, 1]
    mxs, mns = allmxs[:, 0], allmns[:, 0]
    qpx = qPoints[:, 0]
    qpSIdxs, mxSIdxs, mnSIdxs = np.argsort(qpx), np.argsort(mxs), np.argsort(mns)
    qpSIdx, mxSIdx, mnSIdx = 0, 0, 0
    qpIdx, mxIdx, mnIdx = qpSIdxs[qpSIdx], mxSIdxs[mxSIdx], mnSIdxs[mnSIdx]
    qp, mx, mn = qpx[qpIdx], mxs[mxIdx], mns[mnIdx]

    out = {}
    missing = set()

    if pBar is not None:
        pBar.setValue(0)
        allVals = len(qpSIdxs) + len(mxSIdxs) + len(mnSIdxs)
        pBar.setRange(0, allVals)
        pBar.setLabelText("Sweeping ...")
        from Qt.QtWidgets import QApplication

        QApplication.processEvents()

    # skip any triangles to the left of the first query point
    # The algorithm will stop once we hit the last query point
    skip = [qp > m for m in mxs]
    activeTris = [False for _ in mxs]  # For fast membership testing
    atSet = set()  # For fast iteration

    while True:
        if pBar is not None:
            cVal = qpSIdx + mnSIdx + mxSIdx
            if cVal % 1283 == 0:  # Just a random prime
                pBar.setValue(cVal)
                pBar.setLabelText("Sweeping ...\n{0}/{1}".format(cVal, allVals))
                QApplication.processEvents()
                if pBar.wasCanceled():
                    raise RuntimeError("Cancelled!")

        if mn <= mx and mn <= qp:
            # Always add triangles first if possible
            aTriIdx = mnSIdxs[mnSIdx]
            if not activeTris[aTriIdx] and not skip[aTriIdx]:
                activeTris[aTriIdx] = True
                atSet.add(aTriIdx)

            mnSIdx += 1
            if mnSIdx != len(mnSIdxs):
                mnIdx = mnSIdxs[mnSIdx]
                mn = mns[mnIdx]
            else:
                mn = mxs[mxSIdxs[-1]] + 1

        elif qp <= mx:
            # Check query points between adding and removing
            qPoint = qPoints[qpIdx]
            yv = qPoint[1]
            # A linear search is faster than more complex collections here
            for t in atSet:
                if ymns[t] <= yv <= ymxs[t]:
                    a, b, c = uvs[tris[t]]
                    if pointInTri(qPoint, a, b, c):
                        out[qpIdx] = t
                        break
            else:
                missing.add(qpIdx)

            qpSIdx += 1
            if qpSIdx == len(qpSIdxs):
                break
            qpIdx = qpSIdxs[qpSIdx]
            qp = qpx[qpIdx]

        else:
            # always remove triangles last
            aTriIdx = mxSIdxs[mxSIdx]
            if activeTris[aTriIdx] and not skip[aTriIdx]:
                activeTris[aTriIdx] = False
                atSet.remove(aTriIdx)

            mxSIdx += 1
            if mxSIdx != len(mxSIdxs):
                mxIdx = mxSIdxs[mxSIdx]
                mx = mxs[mxIdx]

    return out, missing

########################
# Triangulation
########################

def inBox(point, mxs, mns):
    """Check if a point is inside the bounding box

    Parameters
    ----------
    point : np.array
        A 3d point
    mxs : np.array
        The max values of all the coordinates per axis
    mns : np.array
        The min values of all the coordinates per axis

    Returns
    -------
    : bool
        Whether the point is in the given bounding box

    """
    return np.all(point <= mxs) and np.all(point >= mns)

def _isEar(a, b, c, polygon, tol=EPS):
    """Check if the points a,b,c of the polygon could be their own triangle """
    signedArea = triArea(a, b, c)

    # Check that the triange is wound the correct way.
    if signedArea > 0:
        return False

    # Check that the triangle has non-zero area
    # we already know the area is negative from above
    if -signedArea < tol:
        return False

    # Check that none of the other points in the polygon are contained in triangle
    for p in polygon:
        if p not in (a, b, c):
            if pointInTri(p, a, b, c, tol=tol):
                return False
    return True

def earclip(idxs, verts):
    """Simple earclipping algorithm
        For a polygon with n points it will return n-2 triangles.

    Parameters
    ----------
    idxs : [int, ...]
        The indices that make up a polygon
    verts : np.array
        All the UV positions

    Returns
    -------
    : [[int, int, int], ...]
        A list of triangle indices

    """
    earVerts = []
    tris = []
    idxs = list(idxs)
    polygon = [tuple(verts[i]) for i in idxs]

    numPts = len(polygon)
    for i in range(numPts):
        prev = polygon[i - 2]
        cur = polygon[i - 1]
        nxt = polygon[i]
        if _isEar(prev, cur, nxt, polygon):
            earVerts.append(cur)

    while earVerts and numPts >= 3:
        ear = earVerts.pop()
        i = polygon.index(ear)
        pi, ni = i - 1, (i + 1) % numPts

        prevPt = polygon[pi]
        nxtPt = polygon[ni]
        tris.append((idxs[pi], idxs[i], idxs[ni]))

        polygon.pop(i)
        idxs.pop(i)
        numPts -= 1
        if numPts > 3:
            prePrePt = polygon[i - 2]
            nxtNxtPt = polygon[(i + 1) % numPts]

            groups = [
                (prePrePt, prevPt, nxtPt, polygon),
                (prevPt, nxtPt, nxtNxtPt, polygon),
            ]

            for group in groups:
                p = group[1]
                if _isEar(*group):
                    if p not in earVerts:
                        earVerts.append(p)
                elif p in earVerts:
                    earVerts.remove(p)
    return tris

def triangulateUVs(faces, uvs):
    """Take a set of uvFaces and uv points, and triangulate it

    Parameters
    ----------
    faces : [[int, ...], ...]
        A uv face structure
    uvs : np.array
        The uv positions

    Returns
    -------
    : np.array
        A Nx3 array of uv indexes making triangles
    : [int, ...]
        A list where the index is the triangle index, and the value is the face index
    : {(int, int): int, ...}
        A dictionary of border edge pairs to border faces

    """
    uvs = np.array(uvs)
    triMap = []
    tris = []
    borderFaceMap = {}
    # Build the naiive triangulation
    # And get the borders while I'm looping
    for f, face in enumerate(faces):
        for i in range(2, len(face)):
            triMap.append(f)
            tris.append((face[0], face[i - 1], face[i]))

        for i in range(len(face)):
            ep = (face[i - 1], face[i])
            if ep[0] > ep[1]:
                ep = (ep[1], ep[0])
            if ep in borderFaceMap:
                borderFaceMap.pop(ep)
            else:
                borderFaceMap[ep] = f

    tris = np.array(tris)
    tuvs = uvs[tris]

    # Get the area using a 2d-pseudo-cross-product
    a = tuvs[:, 0] - tuvs[:, 1]
    b = tuvs[:, 2] - tuvs[:, 1]
    signedAreas = a[:, 0] * b[:, 1] - a[:, 1] * b[:, 0]

    # Things with negative area are wound backwards in this case
    negArea = np.where(signedAreas < 0)

    # Do we need to retriangulate any of the polys
    retri = sorted(set(triMap[i] for i in negArea[0]))
    if retri:
        tris = tris.tolist()
        tmpFaceMap = {}
        # tmpFaceMap is only good while we re-triangulate
        # it has to be re-built after
        for t, f in enumerate(triMap):
            tmpFaceMap.setdefault(f, []).append(t)

        for f in retri[::-1]:
            good = earclip(faces[f], uvs)
            tidxs = tmpFaceMap[f]
            tris[tidxs[0]: tidxs[-1] + 1] = good
            triMap[tidxs[0]: tidxs[-1] + 1] = [f] * len(good)
        tris = np.array(tris)

    return tris, triMap, borderFaceMap

########################
# UV Transferring
########################

MISSING = {}

def cooSparseMul(M, v):
    """Multply the sparse matrix by the vector

    Parameters
    ----------
    M : SparseMatrix
        My own sparse matrix representation
    v : np.array
        A vector

    Returns
    -------
    : np.array
        The result of the multiplication
    """
    # Using scipy is at least 2x faster than the
    # pure numpy implementation. But even that
    # is about 5x faster than using a dense matrix
    try:
        from scipy import sparse
    except ImportError:
        # No scipy, use numpy
        # The numpy universal function (ufunc) .add.at()
        # makes this possible without any python looping
        row, col, val, shape = M
        v = v.swapaxes(0, -2)
        vshape = val.shape + tuple([1] * (len(v) - 1))
        val = val.reshape(vshape)
        out = np.zeros(shape[:1] + v.shape[1:])
        np.add.at(out, row, val * v[col])
        out = out.swapaxes(0, -2)
    else:
        row, col, val, shape = M
        spM = sparse.coo_matrix((val, (row, col)), shape=shape)
        spM = sparse.csr_matrix(spM)
        if len(v.shape) == 3:
            out = np.array([spM.dot(frm) for frm in v])
        else:
            out = spM.dot(v)
    return out

def getUvCorrelation(samples, points, faces, tol=0.0001, handleMissing=True, pBar=None):
    """Get the per-face correlation of the sample points in uv space.

    Parameters
    ----------
    samples : np.array
        The sample points
    points : np.array
        The background points
    faces : [[int, ...], ...]
        The Face list
    tol : float
        A small tolerance value, defaulting to the global EPS
    handleMissing :
         (Default value = True)
    pBar : QProgressDialog, optional
        An optional progress dialog

    Returns
    -------
    : {idx: (idx, [float, ...]), ...}
        A dictionary from the sample index to a faceIdx/cornerWeights pair

    """
    tris, triMap, borderMap = triangulateUVs(faces, points)
    swept, missing = sweep(samples, points, tris, pBar=pBar)
    uvToFace = {uvI: triMap[tI] for uvI, tI in six.iteritems(swept)}
    # import __main__
    # __main__.__dict__.update(locals())
    # raise RuntimeError("STOPPIT")
    if pBar is not None:
        pBar.setValue(0)
        pBar.setRange(0, len(uvToFace))
        pBar.setLabelText("Calculate Mean Value Coords")
        from Qt.QtWidgets import QApplication

        QApplication.processEvents()

    # mvcDict = {}
    # for i, (uvIdx, faceIdx) in enumerate(uvToFace.iteritems()):
    # if pBar is not None:
    # pBar.setValue(i)
    # pBar.setLabelText("Calculate Mean Value Coords\n{0}/{1}".format(i, len(uvToFace)))
    # QApplication.processEvents()
    # uv = samples[uvIdx]
    # corners = points[faces[faceIdx]]
    # bary = mvc(corners, uv)
    # mvcDict[uvIdx] = (faceIdx, bary)

    mvcDict = mmvc(faces, points, samples, uvToFace)

    # find the closest border
    # Just use a brute-force search
    # ... for now

    if missing and handleMissing:
        if pBar is not None:
            pBar.setValue(0)
            pBar.setLabelText("Handle Missing")
            QApplication.processEvents()
        tol = tol ** 2
        bk = list(borderMap.keys())
        borders = np.array(bk)
        bStarts = points[borders[:, 0]]
        bEnds = points[borders[:, 1]]

        bDiff = bEnds - bStarts
        bLens2 = (bDiff * bDiff).sum(axis=1)
        for mIdx in missing:
            mp = samples[mIdx]
            dSquared = _mpCheck(bStarts, bDiff, bLens2, mp)
            minIdx = np.argmin(dSquared)
            if dSquared[mIdx] < tol:
                faceIdx = borderMap[bk[minIdx]]
                corners = points[faces[faceIdx]]
                bary = mvc(corners, mp)
                mvcDict[mIdx] = (faceIdx, bary)

    return mvcDict

def _mpCheck(a, d, dr2, pt):
    """Point to Multi-segment squared distance. Uses pre-computed values """
    lerp = ((pt - a) * d).sum(axis=1) / dr2
    lerp = np.clip(lerp, 0, 1)
    xy = (lerp[:, None] * d) + a
    _dxy = xy - pt
    return (_dxy * _dxy).sum(axis=1)

def getVertCorrelation(
    sUvFaces, sUvs, tVertFaces, tUvFaces, tUvs, tol=0.0001, pBar=None
):
    """Build the vertex position correlation between two meshes
        by looking through UV-space. Handle combining multiple uvs
        per vertex, and having samples that live outside the coverage

    Parameters
    ----------
    sUvFaces : [[int, ...], ...]
        The source UVFace list
    sUvs : np.array
        The source UV positions
    tVertFaces : [[int, ...], ...]
        The target vertex face list
    tUvFaces : [[int, ...], ...]
        The target UV face list
    tUvs : np.array
        The target UV positions
    tol : float
        A small tolerance value, defaulting to the global EPS
    pBar : QProgressDialog, optional
        An optional progress dialog

    Returns
    -------
    : {idx: (idx, [float, ...]), ...}
        A dictionary from the sample index to a faceIdx/cornerWeights pair
    """
    childUvToVert = {}
    cNumVerts = -1
    for vF, uvF in zip(tVertFaces, tUvFaces):
        for vIdx, uvIdx in zip(vF, uvF):
            cNumVerts = max(cNumVerts, vIdx)
            childUvToVert[uvIdx] = vIdx

    mvcUvDict = getUvCorrelation(tUvs, sUvs, sUvFaces, tol=tol, pBar=pBar)

    mvcVertDict = {}
    for uvIdx in range(len(tUvs)):
        if uvIdx not in mvcUvDict:
            continue
        mvcVertDict.setdefault(childUvToVert[uvIdx], []).append(mvcUvDict[uvIdx])

    missing = set(range(cNumVerts)) - six.viewkeys(mvcVertDict)
    if missing:
        import time

        v = time.time()
        print(
            "Missing correspondences found. Stored in uvTransfer.MISSING[{0}]".format(v)
        )
        MISSING[v] = missing

    return mvcVertDict

def applyTransfer(parVerts, parFaces, correlation, outputSize):
    """Given a vertex corelation, a driver, and driven points,
        Apply the driver deformation to the driven. This could be
        for one frame, or many frames

    Parameters
    ----------
    parVerts : np.array
        The "parent" vertex positions
    parFaces : [[int, ...], ...]
        The "parent" face list
    correlation : {int: [int, ...], ...}
        A dict correlatng the vertIdx to its possible correlations
    outputSize : (int, ...)
        A numpy output size tuple

    Returns
    -------
    : np.array
        The new vertex positions

    """
    if len(parVerts.shape) == 2:
        parVerts = parVerts[None, ...]

    rows, cols, vals = [], [], []
    for cVertIdx, corrPoss in six.iteritems(correlation):
        if len(corrPoss) == 1:
            pFaceIdx, bary = corrPoss[0]
        else:
            # pick the one with the highest sum-of-squares
            x = [sum(bary * bary) for _, bary in corrPoss]
            idx = x.index(max(x))
            pFaceIdx, bary = corrPoss[idx]

        rows.extend([cVertIdx] * len(bary))
        vals.extend(bary)
        cols.extend(parFaces[pFaceIdx])
    M = np.array(rows), np.array(cols), np.array(vals), (outputSize, parVerts.shape[-2])
    out = cooSparseMul(M, parVerts)
    return out

def uvTransfer(
    srcFaces,
    srcUvFaces,
    srcVerts,
    srcUvs,
    tarFaces,
    tarUvFaces,
    tarVerts,
    tarUvs,
    tol=0.0001,
    pBar=None,
):
    """A helper function that transfers pre-loaded data.
        The source data will be transferred onto the tar data

    Parameters
    ----------
    srcFaces : [[int, ...], ...]
        The source vertex face list
    srcUvFaces : [[int, ...], ...]
        The source uv face list
    srcUvs : np.array
        The source UV positions
    tarFaces : [[int, ...], ...]
        The target vertex face list
    tarUvFaces : [[int, ...], ...]
        The target uv face list
    tarUvs : np.array
        The Target UV Positions
    srcVerts : np.array
        The source Vertex positions
    tarVerts : np.array
        The target Vertex positions
    tol : float
        A small tolerance value, defaulting to the global EPS
    pBar : QProgressDialog, optional
        An optional progress dialog

    Returns
    -------
    : np.array
        The new target vert positions

    """
    corr = getVertCorrelation(
        srcUvFaces, srcUvs, tarFaces, tarUvFaces, tarUvs, tol=tol, pBar=pBar
    )
    if pBar is not None:
        pBar.setValue(0)
        pBar.setLabelText("Apply Transfer")
        from Qt.QtWidgets import QApplication

        QApplication.processEvents()

    return applyTransfer(srcVerts, srcFaces, corr, len(tarVerts))

def uvTransferLoad(
    srcPath, tarPath, srcUvSet="default", tarUvSet="default", tol=0.0001, pBar=None
):
    """Transfer the shape from the source to the target through uv space
        Return the data needed to write out the result

    Parameters
    ----------
    srcPath : str
        The source mesh path (obj, abc, or smpx)
    tarPath : str
        The target mesh path (obj, abc, or smpx)
    srcUvSet : str
        The name of the uv set to use on the source
    tarUvSet : str
        The name of the uv set to use on the target
    tol : float
        A small tolerance value, defaulting to the global EPS
    pBar : QProgressDialog, optional
        An optional progress dialog

    Returns
    -------
    : np.array
        The target vertex positions
    : list
        The target vertex faces
    : np.array
        The target uvs
    : list
        The target uv faces

    """
    from blur3d.api.classes.mesh import Mesh
    from blur3d.api.classes import abc

    if srcPath.endswith(".abc") or srcPath.endswith(".smpx"):
        src = Mesh.loadAbc(srcPath, ensureWinding=False)
        srcVerts = abc.getSampleArray(abc.getMesh(srcPath))
    elif srcPath.endswith(".obj"):
        src = Mesh.loadObj(srcPath, ensureWinding=False)
        srcVerts = np.array(src.vertArray)

    if tarPath.endswith(".abc"):
        tar = Mesh.loadAbc(tarPath, ensureWinding=False)
    elif tarPath.endswith(".obj"):
        tar = Mesh.loadObj(tarPath, ensureWinding=False)

    srcFaces = src.faceVertArray
    srcUvFaces = src.uvFaceMap[srcUvSet]
    srcUvs = np.array(src.uvMap[srcUvSet])

    tarFaces = tar.faceVertArray
    tarUvFaces = tar.uvFaceMap[tarUvSet]
    tarUvs = np.array(tar.uvMap[tarUvSet])
    oldTarVerts = np.array(tar.vertArray)
    tarVerts = uvTransfer(
        srcFaces,
        srcUvFaces,
        srcVerts,
        srcUvs,
        tarFaces,
        tarUvFaces,
        oldTarVerts,
        tarUvs,
        tol=tol,
        pBar=pBar,
    )

    return tarVerts, tarFaces, tarUvs, tarUvFaces

def uvTransferFiles(
    srcPath,
    tarPath,
    outAbcPath,
    srcUvSet="default",
    tarUvSet="default",
    tol=0.0001,
    pBar=None,
):
    """Transfer the shape from the source to the target through uv space
        and write out the result

    Parameters
    ----------
    srcPath : str
        The source mesh path (obj, abc, or smpx)
    tarPath : str
        The target mesh path (obj, abc, or smpx)
    outAbcPath : str
        The path to the output .abc file
    srcUvSet : str
        The name of the uv set to use on the source
    tarUvSet : str
        The name of the uv set to use on the target
    tol : float
        A small tolerance value, defaulting to the global EPS
    pBar : QProgressDialog, optional
        An optional progress dialog

    Returns
    -------

    """
    from blur3d.api.classes import abc

    tarVerts, tarFaces, tarUvs, tarUvFaces = uvTransferLoad(
        srcPath, tarPath, srcUvSet=srcUvSet, tarUvSet=tarUvSet, tol=tol, pBar=pBar
    )
    abc.buildAbc(outAbcPath, tarVerts, tarFaces, uvs=tarUvs, uvFaces=tarUvFaces)


from __future__ import annotations

from typing import cast, Union, TypeVar, Protocol, Optional, Union
from collections.abc import Iterable

from time import perf_counter
import logging

from .utils import isclose

ObjectiveBound = Union[float, int]

Objective = TypeVar('Objective', bound=ObjectiveBound, covariant=True)
LocalMove = TypeVar('LocalMove')

class SolutionProtocol(Protocol[Objective, LocalMove]):
    def random_local_moves_wor(self) -> Iterable[LocalMove]: ...
    def objective_incr_local(self, lmove: LocalMove) -> Optional[Objective]: ...
    def step(self, lmove: LocalMove) -> None: ...

Solution = TypeVar('Solution', bound=SolutionProtocol)

def first_improvement(solution: Solution, budget: float) -> Solution:
    "First-improvement local search"
    start = perf_counter()
    while perf_counter() - start < budget:
        for move in solution.random_local_moves_wor():
            incr = cast(ObjectiveBound, solution.objective_incr_local(move))
            if incr < 0 and not isclose(incr, 0):
                logging.debug(f"Improvement found: {incr}")
                solution.step(move)
                break
            if perf_counter() - start >= budget:
                return solution
        else:
            break
    return solution


# find sum of negative numbers
def neg_sum(list1):
    """
    this method is finds the sum of negative numbers
    in the list that we are passing
    :param list1: input the list
    :return: sum of all negative numbers inside list1
    """
    # counter for sum of
    # negative numbers
    neg_sum1 = 0

    for num in list1:

        # converting number
        # to integer explicitly
        num = int(num)

        # if negative number
        if num < 0:
            # simply add to the
            # negative sum
            neg_sum1 += num
    return neg_sum1

# find sum of positive numbers
# according to categorize
def pos_sum(list1):
    """
    this method is finds the sum of positive numbers
    in the list that we are passing
    :param list1: input the list
    :return: sum of all positive numbers inside list1
    """
    # counter for sum of
    # positive  numbers
    pos_sum1 = 0
    for num in list1:

        # converting number
        # to integer explicitly
        num = int(num)

        # if negative number
        if num >= 0:
            pos_sum1 += num

    return pos_sum1

if __name__ == '__main__':
    # input a list of numbers
    list_num = [-7, 5, 60, -34, 1]

    # creating an object of class

    # calling method for sum
    # of all negative numbers
    print(neg_sum(list_num))

    # calling method for
    # sum of all positive numbers
    print(pos_sum(list_num))


# %%
import os

os.environ["ACCELERATE_DISABLE_RICH"] = "1"
import sys
import pandas as pd
import torch as t
from torch import Tensor, optim
import torch.nn.functional as F
from torchvision import datasets
from torch.utils.data import DataLoader, Subset
from typing import Callable, Iterable, List, Tuple, Optional
import pytorch_lightning as pl
from pytorch_lightning.loggers import CSVLogger, WandbLogger
from dataclasses import dataclass
from pathlib import Path
import numpy as np
from IPython.display import display, HTML

# Make sure exercises are in the path
chapter = r"chapter0_fundamentals"
exercises_dir = Path(f"{os.getcwd().split(chapter)[0]}/{chapter}/exercises").resolve()
section_dir = exercises_dir / "part4_optimization"
if str(exercises_dir) not in sys.path:
    sys.path.append(str(exercises_dir))
os.chdir(section_dir)

from plotly_utils import bar, imshow
from part3_resnets.solutions import (
    IMAGENET_TRANSFORM,
    get_resnet_for_feature_extraction,
    plot_train_loss_and_test_accuracy_from_metrics,
)
from part4_optimization.utils import plot_fn, plot_fn_with_points
import part4_optimization.tests as tests

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

MAIN = __name__ == "__main__"

# %%
import einops

def pathological_curve_loss(x: t.Tensor, y: t.Tensor, angle=0):
    # Example of a pathological curvature. There are many more possible, feel free to experiment here!
    rotations = t.tensor(
        [
            [np.cos(angle), np.sin(angle)],
            [-np.sin(angle), np.cos(angle)],
        ],
        dtype=x.dtype,
    )
    xy = t.stack([x, y], dim=-1)
    x, y = einops.einsum(rotations, xy, "a b, ... b -> a ...")
    x_loss = t.tanh(x) ** 2 + 0.01 * t.abs(x)
    y_loss = t.sigmoid(y)
    return x_loss + y_loss

if MAIN:
    plot_fn(pathological_curve_loss)

# %%
def opt_fn_with_sgd(
    fn: Callable, xy: t.Tensor, lr=0.001, momentum=0.98, n_iters: int = 100
):
    """
    Optimize the a given function starting from the specified point.

    xy: shape (2,). The (x, y) starting point.
    n_iters: number of steps.
    lr, momentum: parameters passed to the torch.optim.SGD optimizer.

    Return: (n_iters, 2). The (x,y) BEFORE each step. So out[0] is the starting point.
    """
    results = []
    optim = t.optim.SGD([xy], lr=lr, momentum=momentum)

    for i in range(n_iters):
        optim.param_groups[0]["momentum"] = momentum ** (i / n_iters)
        results.append(xy.tolist())
        val = fn(*xy)
        val.backward()
        optim.step()
        optim.zero_grad()

    return t.tensor(results)

if MAIN:
    xy = t.tensor([2.5, 2.5], requires_grad=True)
    out = opt_fn_with_sgd(pathological_curve_loss, xy, lr=0.02, momentum=0.99)
    print(out[-1])

# %%
if MAIN:
    points = []

    optimizer_list = [
        (optim.SGD, {"lr": 0.1, "momentum": 0.0}),
        (optim.SGD, {"lr": 0.02, "momentum": 0.99}),
        (optim.SGD, {"lr": 0.06, "momentum": 0.5}),
        (optim.SGD, {"lr": 0.06, "momentum": 0.5}),
    ]
    optimizer_list += [
        (optim.SGD, {"lr": 0.4, "momentum": 0.9 + x / 1000}) for x in range(0, 101, 10)
    ]
    n_iters = 500

    for optimizer_class, params in optimizer_list:
        xy = t.tensor([2.5, 2.5], requires_grad=True)
        xys = opt_fn_with_sgd(
            pathological_curve_loss,
            xy=xy,
            lr=params["lr"],
            momentum=params["momentum"],
            n_iters=n_iters,
        )

        points.append((xys, optimizer_class, params))

    plot_fn_with_points(pathological_curve_loss, points=points[-9:])

for x in points:
    print(x[0][-1])

# %% Plot the evolution of momentum
class SGD:
    def __init__(
        self,
        params: Iterable[t.nn.parameter.Parameter],
        lr: float,
        momentum: float = 0.0,
        weight_decay: float = 0.0,
    ):
        """Implements SGD with momentum.

        Like the PyTorch version, but assume nesterov=False, maximize=False, and dampening=0
            https://pytorch.org/docs/stable/generated/torch.optim.SGD.html#torch.optim.SGD

        """
        self.params = list(
            params
        )  # turn params into a list (because it might be a generator)
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay

        self.running_momentums: List[Optional[Tensor]] = [None] * len(self.params)

    def zero_grad(self) -> None:
        for param in self.params:
            param.grad = None

    @t.inference_mode()
    def step(self) -> None:
        for i, param in enumerate(self.params):
            if self.weight_decay != 0:
                grad = param.grad + self.weight_decay * param
            else:
                grad = param.grad
            assert grad is not None

            if self.momentum != 0:
                if self.running_momentums[i] is None:
                    self.running_momentums[i] = grad  # clone?
                else:
                    self.running_momentums[i].set_(
                        self.momentum * self.running_momentums[i] + grad
                    )
                grad = self.running_momentums[i]
            param.add_(grad, alpha=-self.lr)

    def __repr__(self) -> str:
        return f"SGD(lr={self.lr}, momentum={self.mu}, weight_decay={self.lmda})"

if MAIN:
    tests.test_sgd(SGD)

# %%
class RMSprop:
    def __init__(
        self,
        params: Iterable[t.nn.parameter.Parameter],
        lr: float = 0.01,
        alpha: float = 0.99,
        eps: float = 1e-08,
        weight_decay: float = 0.0,
        momentum: float = 0.0,
    ):
        """Implements RMSprop.

        Like the PyTorch version, but assumes centered=False
            https://pytorch.org/docs/stable/generated/torch.optim.RMSprop.html

        """
        self.params = list(
            params
        )  # turn params into a list (because it might be a generator)
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.alpha = alpha
        self.eps = eps

        self.v = [t.zeros_like(param) for param in self.params]
        self.b = [t.zeros_like(param) for param in self.params]

    def zero_grad(self) -> None:
        for param in self.params:
            param.grad = None

    @t.inference_mode()
    def step(self) -> None:
        for param, v, b in zip(self.params, self.v, self.b):
            grad = param.grad
            if self.weight_decay != 0:
                grad = grad + self.weight_decay * param

            v.set_(self.alpha * v + (1 - self.alpha) * grad.square())

            if self.momentum > 0:
                b.set_(self.momentum * b + grad / (v.sqrt() + self.eps))
                param.add_(b, alpha=-self.lr)
            else:
                param.add_(grad / (v.sqrt() + self.eps), alpha=-self.lr)

    def __repr__(self) -> str:
        return f"RMSprop(lr={self.lr}, eps={self.eps}, momentum={self.mu}, weight_decay={self.lmda}, alpha={self.alpha})"

if MAIN:
    tests.test_rmsprop(RMSprop)

# %%
class Adam:
    def __init__(
        self,
        params: Iterable[t.nn.parameter.Parameter],
        lr: float = 0.001,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-08,
        weight_decay: float = 0.0,
    ):
        """Implements Adam.

        Like the PyTorch version, but assumes amsgrad=False and maximize=False
            https://pytorch.org/docs/stable/generated/torch.optim.Adam.html
        """
        self.params = list(params)
        self.lr = lr
        self.betas = betas
        self.eps = eps
        self.weight_decay = weight_decay

        self.m = [t.zeros_like(param) for param in self.params]
        self.v = [t.zeros_like(param) for param in self.params]

        self.t = 0

    def zero_grad(self) -> None:
        for param in self.params:
            param.grad = None

    @t.inference_mode()
    def step(self) -> None:
        self.t += 1
        for param, m, v in zip(self.params, self.m, self.v):
            grad = param.grad
            assert grad is not None
            if self.weight_decay != 0:
                grad = grad + self.weight_decay * param
            m.set_(self.betas[0] * m + (1 - self.betas[0]) * grad)
            v.set_(self.betas[1] * v + (1 - self.betas[1]) * grad.square())

            mhat = m / (1 - self.betas[0] ** self.t)
            vhat = v / (1 - self.betas[1] ** self.t)

            param.add_(mhat / (vhat.sqrt() + self.eps), alpha=-self.lr)

    def __repr__(self) -> str:
        return f"Adam(lr={self.lr}, beta1={self.beta1}, beta2={self.beta2}, eps={self.eps}, weight_decay={self.lmda})"

if MAIN:
    tests.test_adam(Adam)

# %%
class AdamW:
    def __init__(
        self,
        params: Iterable[t.nn.parameter.Parameter],
        lr: float = 0.001,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-08,
        weight_decay: float = 0.0,
    ):
        """Implements Adam.

        Like the PyTorch version, but assumes amsgrad=False and maximize=False
            https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html
        """
        self.params = list(params)
        self.lr = lr
        self.betas = betas
        self.eps = eps
        self.weight_decay = weight_decay

        self.m = [t.zeros_like(param) for param in self.params]
        self.v = [t.zeros_like(param) for param in self.params]

        self.t = 0

    def zero_grad(self) -> None:
        for param in self.params:
            param.grad = None

    @t.inference_mode()
    def step(self) -> None:
        self.t += 1
        for param, m, v in zip(self.params, self.m, self.v):
            grad = param.grad
            assert grad is not None
            if self.weight_decay != 0:
                param *= 1 - self.lr * self.weight_decay

            m.set_(self.betas[0] * m + (1 - self.betas[0]) * grad)
            v.set_(self.betas[1] * v + (1 - self.betas[1]) * grad.square())

            mhat = m / (1 - self.betas[0] ** self.t)
            vhat = v / (1 - self.betas[1] ** self.t)

            param -= self.lr * mhat / (vhat.sqrt() + self.eps)

    def __repr__(self) -> str:
        return f"AdamW(lr={self.lr}, beta1={self.beta1}, beta2={self.beta2}, eps={self.eps}, weight_decay={self.lmda})"

if MAIN:
    tests.test_adamw(AdamW)
# %%

def opt_fn(
    fn: Callable,
    xy: t.Tensor,
    optimizer_class,
    optimizer_hyperparams: dict,
    n_iters: int = 100,
):
    """Optimize the a given function starting from the specified point.

    optimizer_class: one of the optimizers you've defined, either SGD, RMSprop, or Adam
    optimzer_kwargs: keyword arguments passed to your optimiser (e.g. lr and weight_decay)
    """
    results = []
    optim = optimizer_class([xy], **optimizer_hyperparams)

    for i in range(n_iters):
        results.append(xy.tolist())
        val = fn(*xy)
        val.backward()
        optim.step()
        optim.zero_grad()

    return t.tensor(results)

# %%
if MAIN:
    points = []

    optimizer_list = [
        (SGD, {"lr": 0.03, "momentum": 0.99}),
        (RMSprop, {"lr": 0.03, "alpha": 0.99, "momentum": 0.8}),
        (Adam, {"lr": 0.3, "betas": (0.99, 0.99), "weight_decay": 0.001}),
    ]
    optimizer_list += [(AdamW, optimizer_list[-1][1])]

    for optimizer_class, params in optimizer_list:
        xy = t.tensor([2.5, 2.5], requires_grad=True)
        xys = opt_fn(
            pathological_curve_loss,
            xy=xy,
            optimizer_class=optimizer_class,
            optimizer_hyperparams=params,
            n_iters=300,
        )
        points.append((xys, optimizer_class, params))

    plot_fn_with_points(pathological_curve_loss, points=points)

# %%
def get_cifar(subset: int = 1):
    cifar_trainset = datasets.CIFAR10(
        root="./data", train=True, download=True, transform=IMAGENET_TRANSFORM
    )
    cifar_testset = datasets.CIFAR10(
        root="./data", train=False, download=True, transform=IMAGENET_TRANSFORM
    )
    if subset > 1:
        cifar_trainset = Subset(
            cifar_trainset, indices=range(0, len(cifar_trainset), subset)
        )
        cifar_testset = Subset(
            cifar_testset, indices=range(0, len(cifar_testset), subset)
        )
    return cifar_trainset, cifar_testset

# %%

if MAIN:
    cifar_trainset, cifar_testset = get_cifar()

    imshow(
        cifar_trainset.data[:15],
        facet_col=0,
        facet_col_wrap=5,
        facet_labels=[cifar_trainset.classes[i] for i in cifar_trainset.targets[:15]],
        title="CIFAR-10 images",
        height=600,
    )

# %%
if MAIN:
    cifar_trainset, cifar_testset = get_cifar(subset=1)
    cifar_trainset_small, cifar_testset_small = get_cifar(subset=10)

@dataclass
class ResNetFinetuningArgs:
    batch_size: int = 64
    max_epochs: int = 3
    max_steps: int = 500
    optimizer: t.optim.Optimizer = t.optim.Adam
    learning_rate: float = 1e-3
    log_dir: str = os.getcwd() + "/logs"
    log_name: str = "day5-resnet"
    log_every_n_steps: int = 1
    n_classes: int = 10
    subset: int = 10
    trainset: Optional[datasets.CIFAR10] = None
    testset: Optional[datasets.CIFAR10] = None

    def __post_init__(self):
        if self.trainset is None or self.testset is None:
            self.trainset, self.testset = get_cifar(self.subset)
        self.trainloader = DataLoader(
            self.trainset, shuffle=True, batch_size=self.batch_size
        )
        self.testloader = DataLoader(
            self.testset, shuffle=False, batch_size=self.batch_size
        )
        self.logger = CSVLogger(save_dir=self.log_dir, name=self.log_name)

# %%
class LitResNet(pl.LightningModule):
    def __init__(self, args: ResNetFinetuningArgs):
        super().__init__()
        self.resnet = get_resnet_for_feature_extraction(args.n_classes)
        self.args = args

    def _shared_train_val_step(
        self, batch: Tuple[t.Tensor, t.Tensor]
    ) -> Tuple[t.Tensor, t.Tensor, t.Tensor]:
        """
        Convenience function since train/validation steps are similar.
        """
        imgs, labels = batch
        logits = self.resnet(imgs)
        return logits, labels

    def training_step(
        self, batch: Tuple[t.Tensor, t.Tensor], batch_idx: int
    ) -> t.Tensor:
        """
        Here you compute and return the training loss and some additional metrics for e.g.
        the progress bar or logger.
        """
        logits, labels = self._shared_train_val_step(batch)
        loss = F.cross_entropy(logits, labels)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch: Tuple[t.Tensor, t.Tensor], batch_idx: int) -> None:
        """
        Operates on a single batch of data from the validation set. In this step you might
        generate examples or calculate anything of interest like accuracy.
        """
        logits, labels = self._shared_train_val_step(batch)
        classifications = logits.argmax(dim=1)
        accuracy = t.sum(classifications == labels) / len(classifications)
        self.log("accuracy", accuracy)

    def configure_optimizers(self):
        """
        Choose what optimizers and learning-rate schedulers to use in your optimization.
        """
        optimizer = self.args.optimizer(
            self.resnet.out_layers.parameters(), lr=self.args.learning_rate
        )
        return optimizer

# %%
if MAIN:
    args = ResNetFinetuningArgs(
        trainset=cifar_trainset_small, testset=cifar_testset_small
    )
    model = LitResNet(args)

    trainer = pl.Trainer(
        max_epochs=args.max_epochs,
        max_steps=args.max_steps,
        logger=args.logger,
        log_every_n_steps=args.log_every_n_steps,
    )
    trainer.fit(
        model=model, train_dataloaders=args.trainloader, val_dataloaders=args.testloader
    )

    metrics = pd.read_csv(f"{trainer.logger.log_dir}/metrics.csv")

    plot_train_loss_and_test_accuracy_from_metrics(
        metrics, "Feature extraction with ResNet34"
    )

# %%
def test_resnet_on_random_input(n_inputs: int = 3):
    indices = np.random.choice(len(cifar_trainset), n_inputs).tolist()
    classes = [cifar_trainset.classes[cifar_trainset.targets[i]] for i in indices]
    imgs = cifar_trainset.data[indices]
    with t.inference_mode():
        x = t.stack(list(map(IMAGENET_TRANSFORM, imgs)))
        logits: t.Tensor = model.resnet(x)
    probs = logits.softmax(-1)
    if probs.ndim == 1:
        probs = probs.unsqueeze(0)
    for img, label, prob in zip(imgs, classes, probs):
        display(HTML(f"
Classification probabilities (true class = {label})
"))
        imshow(
            img,
            width=200,
            height=200,
            margin=0,
            xaxis_visible=False,
            yaxis_visible=False,
        )
        bar(
            prob,
            x=cifar_trainset.classes,
            template="ggplot2",
            width=600,
            height=400,
            labels={"x": "Classification", "y": "Probability"},
            text_auto=".2f",
            showlegend=False,
        )

if MAIN:
    test_resnet_on_random_input()
# %%
import wandb

@dataclass
class ResNetFinetuningArgsWandb(ResNetFinetuningArgs):
    use_wandb: bool = True
    run_name: Optional[str] = None

    def __post_init__(self):
        super().__post_init__()
        if self.use_wandb:
            self.logger = WandbLogger(
                save_dir=self.log_dir, project=self.log_name, name=self.run_name
            )

# %%
if MAIN:
    args = ResNetFinetuningArgsWandb(
        trainset=cifar_trainset_small, testset=cifar_testset_small
    )
    model = LitResNet(args)

    trainer = pl.Trainer(
        max_epochs=args.max_epochs,
        max_steps=args.max_steps,
        logger=args.logger,
        log_every_n_steps=args.log_every_n_steps,
    )
    trainer.fit(
        model=model, train_dataloaders=args.trainloader, val_dataloaders=args.testloader
    )
    wandb.finish()
# %%
if MAIN:
    # YOUR CODE HERE - fill `sweep_config`
    sweep_config = dict(
        method="random",
        metric=dict(name="accuracy", goal="maximize"),
        parameters=dict(
            batch_size=dict(values=[32, 64, 128, 256]),
            max_epochs=dict(min=1, max=4),
            learning_rate=dict(max=0.1, min=0.0001, distribution="log_uniform_values"),
        ),
    )
    # FLAT SOLUTION END

    tests.test_sweep_config(sweep_config)
# %%
def train() -> None:
    args = ResNetFinetuningArgsWandb(
        trainset=cifar_trainset_small, testset=cifar_testset_small
    )
    args.batch_size = wandb.config["batch_size"]
    args.max_epochs = wandb.config["max_epochs"]
    args.learning_rate = wandb.config["learning_rate"]

    model = LitResNet(args)

    trainer = pl.Trainer(
        max_epochs=args.max_epochs,
        max_steps=args.max_steps,
        logger=args.logger,
        log_every_n_steps=args.log_every_n_steps,
    )
    trainer.fit(
        model=model, train_dataloaders=args.trainloader, val_dataloaders=args.testloader
    )
    wandb.finish()

# %%
if MAIN:
    wandb.init()
    sweep_id = wandb.sweep(sweep=sweep_config, project='day4-resnet-sweep')
    wandb.agent(sweep_id=sweep_id, function=train, count=3)
# %%


# Read all data from file
data = open("input.txt", "r").read()

# Split data into list of food items caried by each elf
elfInventories = [
    [
        int(cals)
        for cals in elf.split("\n")
        if cals != ""
    ]
    for elf in data.split("\n\n")
]

# Find the top 3 elfs carrying the greatest sums of calories
mostCalories = [0, 0, 0]
for elf in elfInventories:
    if sum(elf) > mostCalories[2]:
        if sum(elf) > mostCalories[1]:
            if sum(elf) > mostCalories[0]:
                mostCalories[2] = mostCalories[1]
                mostCalories[1] = mostCalories[0]
                mostCalories[0] = sum(elf)
            else:
                mostCalories[2] = mostCalories[1]
                mostCalories[1] = sum(elf)
        else:
            mostCalories[2] = sum(elf)

print("The most calories carried by an elf is: ", mostCalories[0])
print("The top three elfs are carrying ", sum(mostCalories), " calories")


# -*- coding: utf-8 -*-
"""
@author: t.maybour
"""

import os
import h5py
import numpy as np

import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt

def main():

    plots = cplots()

    c = np.float64(299792458) # speed of light
    epsilon0 = np.float(8.85418782e-12)
    mu0 = np.float(1.25663706e-6)
    nu = np.sqrt(mu0 / epsilon0)

#########################################################################

    with h5py.File('fdtd_data.h5', "r") as hdf:

        ################### Inputs ###################
        bw_nopo = np.uint32(hdf.get('bw_nopo'))
        tv_nopo_input = np.uint32(hdf.get('tv_nopo_input'))
        xv_nopo_input = np.uint32(hdf.get('xv_nopo_input'))
        res = np.uint32(hdf.get('res'))
        ##############################################
        l0 = np.float64(hdf.get('l0'))
        pfbw = np.float64(hdf.get('pfbw'))
        ##############################################
        glen = np.float64(hdf.get('glen'))
        gstart = np.float64(hdf.get('gstart'))
        dn = np.float64(hdf.get('dn'))
        ds = np.float64(hdf.get('ds'))
        gramp_scale = np.float64(hdf.get('gramp_scale'))
        ##############################################
        tlen = np.float64(hdf.get('tlen'))
        mu = np.float64(hdf.get('mu'))
        tstart = np.float64(hdf.get('tstart'))
        delta_t = np.float64(hdf.get('delta_t'))
        ton = np.float64(hdf.get('ton'))
        ##############################################

        ################### DATA #####################
        tv_nopo = np.uint32(hdf.get('tv_nopo'))
        xv_nopo = np.uint32(hdf.get('xv_nopo'))
        n0 = np.float64(hdf.get('n0'))
        gfstart = np.float64(hdf.get('gfstart'))
        gfend = np.float64(hdf.get('gfend'))
        gend = np.float64(hdf.get('gend'))
        ##############################################
        kv = np.array(hdf.get('kv'), dtype=np.float64)
        fk = np.array(hdf.get('fk'), dtype=np.float64)
        xv = np.array(hdf.get('xv'), dtype=np.float64)
        apod = np.array(hdf.get('apod'), dtype=np.float64)
        gm = np.array(hdf.get('gm'), dtype=np.float64)
        tv = np.array(hdf.get('tv'), dtype=np.float64)
        twindow = np.array(hdf.get('twindow'), dtype=np.float64)
        source_real = np.array(hdf.get('source_real'), dtype=np.float64)
        source_imag = np.array(hdf.get('source_imag'), dtype=np.float64)
        efield_mat_real = np.array(hdf.get('efield_mat_real'), dtype=np.float64)
        efield_mat_imag = np.array(hdf.get('efield_mat_imag'), dtype=np.float64)
        hfield_mat_real = np.array(hdf.get('hfield_mat_real'), dtype=np.float64)
        hfield_mat_imag = np.array(hdf.get('hfield_mat_imag'), dtype=np.float64)
        ##############################################

#########################################################################

    efield_mat = np.zeros((tv_nopo,xv_nopo), dtype=np.complex128)
    hfield_mat = np.zeros((tv_nopo,xv_nopo), dtype=np.complex128)

    for ii in range(0, tv_nopo):
        for jj in range(0, xv_nopo):
            efield_mat[ii,jj] = efield_mat_real[ii*xv_nopo+jj] + 1j*efield_mat_imag[ii*xv_nopo+jj]
            hfield_mat[ii,jj] = hfield_mat_real[ii*xv_nopo+jj] + 1j*hfield_mat_imag[ii*xv_nopo+jj]

    edensity = np.zeros((tv_nopo, xv_nopo), dtype=np.float64)
    edensity[:,:] = np.real((n0**2)*np.conj(efield_mat[:,:])*(efield_mat[:,:]) + (c**2)*np.conj(hfield_mat[:,:])*(hfield_mat[:,:]))
    edensity[:,:] = edensity[:,:] / np.max(edensity[:,:])

    source = np.zeros(tv_nopo, dtype=np.complex128)
    source[:] = source_real[:] + 1j*source_imag[:]

#######################################################################
############################## Graphs #################################
#######################################################################

    plots.x_label = 'Wavenumber (1/m)'
    plots.line_plot(kv, fk/np.max(fk), filename='fk')

    plots.x_scale = 1e3
    plots.x_label = 'Length (mm)'

    plots.line_plot(xv, apod, filename='apod')
    plots.line_plot(xv, gm, filename='gm')

    plots.x_scale = 1e12
    plots.x_label = 'Time (ps)'

    plots.line_plot(tv, twindow, filename='twindow')
    plots.line_plot(tv, np.abs(source)**2/np.max(np.abs(source)**2), filename='source')

########################################################################

    plots.width = 32
    plots.height = 16
    plots.axes_label_size = 50
    plots.tick_label_size = 45
    plots.major_tick_length = 8
    plots.major_tick_width = 3

    plots.x_scale = 1e12
    plots.x_label = 'Time (ps)'
    plots.y_scale = 1e3
    plots.y_label = 'Position (mm)'

    plots.z_label = "Energy Density (a.u.)"

    plots.heat_map(tv, xv, np.transpose(edensity), filename='edensity_plot')

#################################################################################
#################################################################################
#################################################################################

class cplots():

    def __init__(self):
        self.filename = ''
        self.save_dir = ''
        self.width = 20
        self.height = 16
        self.axes_label_size = 35
        self.tick_label_size = 35
        self.axes_label_padding = 35
        self.margins = False
        self.x_label = ''
        self.y_label = ''
        self.z_label = ''
        self.x_scale = 1
        self.y_scale = 1
        self.z_scale = 1
        self.linewidth = 6
        self.major_tick_length = 5
        self.major_tick_width = 3
        self.tick_label_padding = 15

    def set_save_dir(self):
        if (len(self.save_dir) > 0) :
            if (self.save_dir[-1] != '/'):
                self.save_dir = self.save_dir + '/'
        else:
            self.save_dir = 'graphs/'
        self.create_dir(self.save_dir)

    def create_dir(self, dir_in):
        if (dir_in[-1] != '/'):
            dir_in = dir_in + '/'
        if not os.path.exists(dir_in):
            os.makedirs(dir_in)

    def axis_setup(self):
        plt.rc('axes', labelsize=self.axes_label_size)
        plt.rc('xtick', labelsize=self.tick_label_size)
        plt.rc('ytick', labelsize=self.tick_label_size)

    def set_kwargs(self, kwargs):
        self.filename = kwargs.get('filename', self.filename)
        self.save_dir = kwargs.get('save_dir', self.save_dir)

    def line_plot(self, x, data, *args, **kwargs):

        self.set_kwargs(kwargs)
        self.set_save_dir()
        self.axis_setup()

        y = np.array(data)

        plt.figure(figsize=(self.width, self.height))
        plt.ticklabel_format(useOffset=False)
        if self.margins == False:
            plt.margins(x=0)
        plt.xlabel(self.x_label, labelpad=self.axes_label_padding)
        plt.ylabel(self.y_label, labelpad=self.axes_label_padding)

        plt.plot(self.x_scale*x,self.y_scale*y,linewidth=self.linewidth)

        plt.tight_layout()
        plt.savefig(self.save_dir + self.filename + ".png")
        plt.close()

    def heat_map(self, x, y, z, *args, **kwargs):

        self.set_kwargs(kwargs)
        self.set_save_dir()
        self.axis_setup()

        u, v = np.meshgrid(x, y)

        fig, ax1 = plt.subplots(constrained_layout=True)

        ax1.tick_params(direction='out', length=self.major_tick_length, width=self.major_tick_width, labelsize=self.axes_label_size, pad=self.tick_label_padding)
        im = ax1.pcolormesh(self.x_scale*u, self.y_scale*v, self.z_scale*z, cmap='viridis')

        ax1.set_xlabel(self.x_label, labelpad=self.axes_label_padding)
        ax1.set_ylabel(self.y_label, labelpad=self.axes_label_padding)

        cbar = fig.colorbar(im, ax=ax1) #need a colorbar to show the intensity scale
        cbar.set_label(self.z_label, labelpad=self.axes_label_padding)
        cbar.ax.tick_params(labelsize=self.axes_label_size)

        #################################################
        w=self.width
        h=self.height
        ax=plt.gca()
        l = ax.figure.subplotpars.left
        r = ax.figure.subplotpars.right
        t = ax.figure.subplotpars.top
        b = ax.figure.subplotpars.bottom
        figw = float(w)/(r-l)
        figh = float(h)/(t-b)
        ax.figure.set_size_inches(figw, figh)
        #################################################

        plt.savefig(self.save_dir + self.filename + ".png")
        plt.close()

#################################################################################
#################################################################################
#################################################################################

if __name__ == '__main__':
    main()


from typing import Dict, List
import numpy as np
from transformers import BatchEncoding
from dataclasses import dataclass
from transformers import AutoTokenizer
import torch
import math
from torch.optim import Optimizer
from typing import Iterable, Tuple
from torch import nn
import random
import string

@dataclass
class DataCollatorForT5MLM:
    """
    [Copied from https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_t5_mlm_flax.py]
    Data collator used for T5 span-masked language modeling.
    It is made sure that after masking the inputs are of length `data_args.max_seq_length` and targets are also of fixed length.
    For more information on how T5 span-masked language modeling works, one can take a look
    at the `official paper <https://arxiv.org/pdf/1910.10683.pdf>`__
    or the `official code for preprocessing <https://github.com/google-research/text-to-text-transfer-transformer/blob/master/t5/data/preprocessors.py>`__ .
    Args:
        tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
            The tokenizer used for encoding the data.
        noise_density (:obj:`float`):
            The probability with which to (randomly) mask tokens in the input.
        mean_noise_span_length (:obj:`float`):
            The average span length of the masked tokens.
        input_length (:obj:`int`):
            The expected input length after masking.
        target_length (:obj:`int`):
            The expected target length after masking.
        pad_token_id: (:obj:`int`):
            The pad token id of the model
        decoder_start_token_id: (:obj:`int):
            The decoder start token id of the model
    """

    tokenizer: AutoTokenizer
    noise_density: float
    mean_noise_span_length: float
    input_length: int
    target_length: int
    pad_token_id: int

    def __call__(self, examples: List[Dict[str, np.ndarray]]) -> BatchEncoding:
        # convert list to dict and tensorize input
        batch = BatchEncoding(
            {
                k: np.array([examples[i][k] for i in range(len(examples))])
                for k, v in examples[0].items()
            }
        )

        input_ids = batch["input_ids"]
        batch_size, expandend_input_length = input_ids.shape

        mask_indices = np.asarray(
            [
                self.random_spans_noise_mask(expandend_input_length)
                for i in range(batch_size)
            ]
        )
        labels_mask = ~mask_indices

        input_ids_sentinel = self.create_sentinel_ids(mask_indices.astype(np.int8))
        labels_sentinel = self.create_sentinel_ids(labels_mask.astype(np.int8))

        batch["input_ids"] = self.filter_input_ids(input_ids, input_ids_sentinel)
        batch["labels"] = self.filter_input_ids(input_ids, labels_sentinel)

        if batch["input_ids"].shape[-1] != self.input_length:
            raise ValueError(
                f"`input_ids` are incorrectly preprocessed. `input_ids` length is {batch['input_ids'].shape[-1]}, but"
                f" should be {self.input_length}."
            )

        if batch["labels"].shape[-1] != self.target_length:
            raise ValueError(
                f"`labels` are incorrectly preprocessed. `labels` length is {batch['labels'].shape[-1]}, but should be"
                f" {self.target_length}."
            )

        batch = {k: torch.from_numpy(v) for k, v in batch.items()}
        return batch

    def create_sentinel_ids(self, mask_indices):
        """
        Sentinel ids creation given the indices that should be masked.
        The start indices of each mask are replaced by the sentinel ids in increasing
        order. Consecutive mask indices to be deleted are replaced with `-1`.
        """
        start_indices = mask_indices - np.roll(mask_indices, 1, axis=-1) * mask_indices
        start_indices[:, 0] = mask_indices[:, 0]

        sentinel_ids = np.where(
            start_indices != 0, np.cumsum(start_indices, axis=-1), start_indices
        )
        sentinel_ids = np.where(
            sentinel_ids != 0, (len(self.tokenizer) - sentinel_ids), 0
        )
        sentinel_ids -= mask_indices - start_indices

        return sentinel_ids

    def filter_input_ids(self, input_ids, sentinel_ids):
        """
        Puts sentinel mask on `input_ids` and fuse consecutive mask tokens into a single mask token by deleting.
        This will reduce the sequence length from `expanded_inputs_length` to `input_length`.
        """
        batch_size = input_ids.shape[0]

        input_ids_full = np.where(sentinel_ids != 0, sentinel_ids, input_ids)
        # input_ids tokens and sentinel tokens are >= 0, tokens < 0 are
        # masked tokens coming after sentinel tokens and should be removed
        input_ids = input_ids_full[input_ids_full >= 0].reshape((batch_size, -1))
        input_ids = np.concatenate(
            [
                input_ids,
                np.full((batch_size, 1), self.tokenizer.eos_token_id, dtype=np.int32),
            ],
            axis=-1,
        )
        return input_ids

    def random_spans_noise_mask(self, length):
        """This function is copy of `random_spans_helper <https://github.com/google-research/text-to-text-transfer-transformer/blob/84f8bcc14b5f2c03de51bd3587609ba8f6bbd1cd/t5/data/preprocessors.py#L2682>`__ .

        Noise mask consisting of random spans of noise tokens.
        The number of noise tokens and the number of noise spans and non-noise spans
        are determined deterministically as follows:
        num_noise_tokens = round(length * noise_density)
        num_nonnoise_spans = num_noise_spans = round(num_noise_tokens / mean_noise_span_length)
        Spans alternate between non-noise and noise, beginning with non-noise.
        Subject to the above restrictions, all masks are equally likely.

        Args:
            length: an int32 scalar (length of the incoming token sequence)
            noise_density: a float - approximate density of output mask
            mean_noise_span_length: a number

        Returns:
            a boolean tensor with shape [length]
        """

        orig_length = length

        num_noise_tokens = int(np.round(length * self.noise_density))
        # avoid degeneracy by ensuring positive numbers of noise and nonnoise tokens.
        num_noise_tokens = min(max(num_noise_tokens, 1), length - 1)
        num_noise_spans = int(np.round(num_noise_tokens / self.mean_noise_span_length))

        # avoid degeneracy by ensuring positive number of noise spans
        num_noise_spans = max(num_noise_spans, 1)
        num_nonnoise_tokens = length - num_noise_tokens

        # pick the lengths of the noise spans and the non-noise spans
        def _random_segmentation(num_items, num_segments):
            """Partition a sequence of items randomly into non-empty segments.
            Args:
                num_items: an integer scalar > 0
                num_segments: an integer scalar in [1, num_items]
            Returns:
                a Tensor with shape [num_segments] containing positive integers that add
                up to num_items
            """
            mask_indices = np.arange(num_items - 1) < (num_segments - 1)
            np.random.shuffle(mask_indices)
            first_in_segment = np.pad(mask_indices, [[1, 0]])
            segment_id = np.cumsum(first_in_segment)
            # count length of sub segments assuming that list is sorted
            _, segment_length = np.unique(segment_id, return_counts=True)
            return segment_length

        noise_span_lengths = _random_segmentation(num_noise_tokens, num_noise_spans)
        nonnoise_span_lengths = _random_segmentation(
            num_nonnoise_tokens, num_noise_spans
        )

        interleaved_span_lengths = np.reshape(
            np.stack([nonnoise_span_lengths, noise_span_lengths], axis=1),
            [num_noise_spans * 2],
        )
        span_starts = np.cumsum(interleaved_span_lengths)[:-1]
        span_start_indicator = np.zeros((length,), dtype=np.int8)
        span_start_indicator[span_starts] = True
        span_num = np.cumsum(span_start_indicator)
        is_noise = np.equal(span_num % 2, 1)

        return is_noise[:orig_length]

def compute_input_and_target_lengths(inputs_length, noise_density, mean_noise_span_length):
    """This function is copy of `random_spans_helper <https://github.com/google-research/text-to-text-transfer-transformer/blob/84f8bcc14b5f2c03de51bd3587609ba8f6bbd1cd/t5/data/preprocessors.py#L2466>`__ .

    [Copied from https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_t5_mlm_flax.py]
    Training parameters to avoid padding with random_spans_noise_mask.
    When training a model with random_spans_noise_mask, we would like to set the other
    training hyperparmeters in a way that avoids padding.
    This function helps us compute these hyperparameters.
    We assume that each noise span in the input is replaced by extra_tokens_per_span_inputs sentinel tokens,
    and each non-noise span in the targets is replaced by extra_tokens_per_span_targets sentinel tokens.
    This function tells us the required number of tokens in the raw example (for split_tokens())
    as well as the length of the encoded targets. Note that this function assumes
    the inputs and targets will have EOS appended and includes that in the reported length.

    Args:
        inputs_length: an integer - desired length of the tokenized inputs sequence
        noise_density: a float
        mean_noise_span_length: a float
    Returns:
        tokens_length: length of original text in tokens
        targets_length: an integer - length in tokens of encoded targets sequence
    """

    def _tokens_length_to_inputs_length_targets_length(tokens_length):
        num_noise_tokens = int(round(tokens_length * noise_density))
        num_nonnoise_tokens = tokens_length - num_noise_tokens
        num_noise_spans = int(round(num_noise_tokens / mean_noise_span_length))
        # inputs contain all nonnoise tokens, sentinels for all noise spans
        # and one EOS token.
        _input_length = num_nonnoise_tokens + num_noise_spans + 1
        _output_length = num_noise_tokens + num_noise_spans + 1
        return _input_length, _output_length

    tokens_length = inputs_length

    while _tokens_length_to_inputs_length_targets_length(tokens_length + 1)[0] <= inputs_length:
        tokens_length += 1

    inputs_length, targets_length = _tokens_length_to_inputs_length_targets_length(tokens_length)

    # minor hack to get the targets length to be equal to inputs length
    # which is more likely to have been set to a nice round number.
    if noise_density == 0.5 and targets_length > inputs_length:
        tokens_length -= 1
        targets_length -= 1
    return tokens_length, targets_length

class AdamWScale(Optimizer):
    """
    This AdamW implementation is copied from Huggingface.
    We modified it with Adagrad scaling by rms of a weight tensor

    Implements Adam algorithm with weight decay fix as introduced in [Decoupled Weight Decay
    Regularization](https://arxiv.org/abs/1711.05101).

    Parameters:
        params (`Iterable[nn.parameter.Parameter]`):
            Iterable of parameters to optimize or dictionaries defining parameter groups.
        lr (`float`, *optional*, defaults to 1e-3):
            The learning rate to use.
        betas (`Tuple[float,float]`, *optional*, defaults to (0.9, 0.999)):
            Adam's betas parameters (b1, b2).
        eps (`float`, *optional*, defaults to 1e-6):
            Adam's epsilon for numerical stability.
        weight_decay (`float`, *optional*, defaults to 0):
            Decoupled weight decay to apply.
        correct_bias (`bool`, *optional*, defaults to `True`):
            Whether or not to correct bias in Adam (for instance, in Bert TF repository they use `False`).
        no_deprecation_warning (`bool`, *optional*, defaults to `False`):
            A flag used to disable the deprecation warning (set to `True` to disable the warning).
    """

    def __init__(
        self,
        params: Iterable[nn.parameter.Parameter],
        lr: float = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-6,
        weight_decay: float = 0.0,
        correct_bias: bool = True,
    ):
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)")
        if not 0.0 <= eps:
            raise ValueError(f"Invalid epsilon value: {eps} - should be >= 0.0")
        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, correct_bias=correct_bias)
        super().__init__(params, defaults)

    @staticmethod
    def _rms(tensor):
        return tensor.norm(2) / (tensor.numel() ** 0.5)

    def step(self, closure=None):
        """
        Performs a single optimization step.

        Arguments:
            closure (`Callable`, *optional*): A closure that reevaluates the model and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead")

                state = self.state[p]
                beta1, beta2 = group["betas"]

                # State initialization
                if len(state) == 0:
                    state["step"] = 0
                    # Exponential moving average of gradient values
                    state["exp_avg"] = torch.zeros_like(p.data)
                    # Exponential moving average of squared gradient values
                    state["exp_avg_sq"] = torch.zeros_like(p.data)

                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]

                state["step"] += 1

                # Decay the first and second moment running average coefficient
                # In-place operations to update the averages at the same time
                exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                denom = exp_avg_sq.sqrt().add_(group["eps"])

                step_size = group["lr"]
                if group["correct_bias"]:  # No bias correction for Bert
                    bias_correction1 = 1.0 - beta1 ** state["step"]
                    bias_correction2 = 1.0 - beta2 ** state["step"]
                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1

                # /Adapt Step from Adafactor
                step_size = step_size * max(1e-3, self._rms(p.data))
                # /Adapt Step from Adafactor

                p.data.addcdiv_(exp_avg, denom, value=-step_size)

                # Just adding the square of the weights to the loss function is *not*
                # the correct way of using L2 regularization/weight decay with Adam,
                # since that will interact with the m and v parameters in strange ways.
                #
                # Instead we want to decay the weights in a manner that doesn't interact
                # with the m/v parameters. This is equivalent to adding the square
                # of the weights to the loss with plain (non-momentum) SGD.
                # Add weight decay at the end (fixed version)
                if group["weight_decay"] > 0.0:
                    p.data.add_(p.data, alpha=(-group["lr"] * group["weight_decay"]))

        return loss

def tokenize_function(examples, tokenizer, in_length):
    tokenizer_out = tokenizer(
        text=examples["text"],
        return_attention_mask=False,
    )

    input_ids = tokenizer_out["input_ids"]

    concatenated_ids = np.concatenate(input_ids)

    total_length = concatenated_ids.shape[0]
    total_length = (total_length // in_length) * in_length

    concatenated_ids = concatenated_ids[:total_length].reshape(-1, in_length)
    result = {"input_ids": concatenated_ids}

    return result

from transformers.data.data_collator import *
@dataclass
class DataCollatorForNI:
    tokenizer: PreTrainedTokenizerBase
    padding: Union[bool, str, PaddingStrategy] = True
    max_source_length: Optional[int] = None
    max_target_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    label_pad_token_id: int = -100
    return_tensors: str = "pt"
    add_task_name: bool = False
    add_task_definition: bool = True
    num_pos_examples: int = 0
    num_neg_examples: int = 0
    add_explanation: bool = False
    tk_instruct: bool = False
    text_only: bool = False

    def __call__(self, batch, return_tensors=None):

        if return_tensors is None:
            return_tensors = self.return_tensors

        sources = []
        for instance in batch:
            if self.tk_instruct:
                all_valid_encodings = [
                    # instruction only
                    {
                        "add_task_name": False,
                        "add_task_definition": True,
                        "num_pos_examples": 0,
                        "num_neg_examples": 0,
                        "add_explanation": False,
                    },
                    # example only
                    {
                        "add_task_name": False,
                        "add_task_definition": False,
                        "num_pos_examples": 2,
                        "num_neg_examples": 0,
                        "add_explanation": False,
                    },
                    # instruction + pos examples
                    {
                        "add_task_name": False,
                        "add_task_definition": True,
                        "num_pos_examples": 2,
                        "num_neg_examples": 0,
                        "add_explanation": False,
                    },
                    # instruction + pos examples + neg examples
                    {
                        "add_task_name": False,
                        "add_task_definition": True,
                        "num_pos_examples": 2,
                        "num_neg_examples": 2,
                        "add_explanation": False,
                    },
                    # instruction + pos (w. explanation)
                    {
                        "add_task_name": False,
                        "add_task_definition": True,
                        "num_pos_examples": 2,
                        "num_neg_examples": 0,
                        "add_explanation": True,
                    },
                ]
                encoding_schema = random.choice(all_valid_encodings)
                add_task_name = encoding_schema["add_task_name"]
                add_task_definition = encoding_schema["add_task_definition"]
                num_pos_examples = encoding_schema["num_pos_examples"]
                num_neg_examples = encoding_schema["num_neg_examples"]
                add_explanation = encoding_schema["add_explanation"]
            else:
                add_task_name = self.add_task_name
                add_task_definition = self.add_task_definition
                num_pos_examples = self.num_pos_examples
                num_neg_examples = self.num_neg_examples
                add_explanation = self.add_explanation

            task_input = ""
            # add the input first.
            task_input += "Now complete the following example -\n"
            task_input += f"Input: {instance['Instance']['input'].strip()}"
            if not task_input[-1] in string.punctuation:
                task_input += "."
            task_input += "\n"
            task_input += "Output: "

            task_name = ""
            if add_task_name:
                task_name += instance["Task"] + ". "

            definition = ""
            if add_task_definition:
                if isinstance(instance["Definition"], list):
                    definition = (
                        "Definition: " + instance["Definition"][0].strip()
                    )
                else:
                    definition = "Definition: " + instance["Definition"].strip()
                if not definition[-1] in string.punctuation:
                    definition += "."
                definition += "\n\n"

            # try to add positive examples.
            pos_examples = []
            for idx, pos_example in enumerate(
                instance["Positive Examples"][:num_pos_examples]
            ):
                pos_example_str = f" Positive Example {idx+1} -\n"
                pos_example_str += f"Input: {pos_example['input'].strip()}"
                if not pos_example_str[-1] in string.punctuation:
                    pos_example_str += "."
                pos_example_str += "\n"
                pos_example_str += f" Output: {pos_example['output'].strip()}"
                if not pos_example_str[-1] in string.punctuation:
                    pos_example_str += "."
                pos_example_str += "\n"
                if add_explanation and "explanation" in pos_example:
                    pos_example_str += (
                        f" Explanation: {pos_example['explanation'].strip()}"
                    )
                    if not pos_example_str[-1] in string.punctuation:
                        pos_example_str += "."
                    pos_example_str += "\n"
                pos_example_str += "\n"
                if (
                    len(
                        self.tokenizer(
                            definition
                            + " ".join(pos_examples)
                            + pos_example_str
                            + task_input
                        )["input_ids"]
                    )
                    <= self.max_source_length
                ):
                    pos_examples.append(pos_example_str)
                else:
                    break

            # try to add negative examples.
            neg_examples = []
            for idx, neg_example in enumerate(
                instance["Negative Examples"][:num_neg_examples]
            ):
                neg_example_str = f" Negative Example {idx+1} -\n"
                neg_example_str += f"Input: {neg_example['input'].strip()}"
                if not neg_example_str[-1] in string.punctuation:
                    neg_example_str += "."
                neg_example_str += "\n"
                neg_example_str += f" Output: {neg_example['output'].strip()}"
                if not neg_example_str[-1] in string.punctuation:
                    neg_example_str += "."
                neg_example_str += "\n"
                if add_explanation and "explanation" in neg_example:
                    neg_example_str += (
                        f" Explanation: {neg_example['explanation'].strip()}"
                    )
                    if not neg_example_str[-1] in string.punctuation:
                        neg_example_str += "."
                    neg_example_str += "\n"
                neg_example_str += "\n"
                if (
                    len(
                        self.tokenizer(
                            definition
                            + " ".join(pos_examples)
                            + " ".join(neg_examples)
                            + neg_example_str
                            + task_input
                        )["input_ids"]
                    )
                    <= self.max_source_length
                ):
                    neg_examples.append(neg_example_str)
                else:
                    break

            source = (
                task_name
                + definition
                + "".join(pos_examples)
                + "".join(neg_examples)
                + task_input
            )
            tokenized_source = self.tokenizer(source)["input_ids"]
            if len(tokenized_source) <= self.max_source_length:
                sources.append(source)
            else:
                sources.append(
                    self.tokenizer.decode(
                        tokenized_source[: self.max_source_length],
                        skip_special_tokens=True,
                    )
                )

        if self.text_only:
            model_inputs = {"inputs": sources}
        else:
            model_inputs = self.tokenizer(
                sources,
                max_length=self.max_source_length,
                padding=self.padding,
                return_tensors=self.return_tensors,
                truncation=True,
                pad_to_multiple_of=self.pad_to_multiple_of,
            )

        if "output" in batch[0]["Instance"] and batch[0]["Instance"]["output"]:
            # Randomly select one reference if multiple are provided.
            labels = [random.choice(ex["Instance"]["output"]) for ex in batch]
            if self.text_only:
                model_inputs["labels"] = labels
            else:
                labels = self.tokenizer(
                    labels,
                    max_length=self.max_target_length,
                    padding=self.padding,
                    return_tensors=self.return_tensors,
                    truncation=True,
                    pad_to_multiple_of=self.pad_to_multiple_of,
                )
                label_mask = labels["attention_mask"].bool()
                model_inputs["labels"] = labels["input_ids"].masked_fill(
                    ~label_mask, self.label_pad_token_id
                )
        else:
            model_inputs["labels"] = None

        return model_inputs


#  This file is part of cappuccino.
#
#  cappuccino is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  cappuccino is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with cappuccino.  If not, see <https://www.gnu.org/licenses/>.

import irc3
import sentry_sdk
from irc3.plugins.command import command
from requests import RequestException
from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration

from cappuccino import Plugin

def _before_send(event, hint):
    if "exc_info" in hint:
        exc_type, exc_value, tb = hint["exc_info"]
        if isinstance(exc_value, (RequestException, TimeoutError)):
            return None

    return event

@irc3.plugin
class Sentry(Plugin):
    requires = ["irc3.plugins.command", "cappuccino.core"]

    def __init__(self, bot):
        super().__init__(bot)

        dsn = self.config.get("dsn", None)
        if not dsn:
            self.logger.warning("Missing Sentry DSN, Sentry will not be used.")
            return

        sentry_sdk.init(
            dsn,
            before_send=_before_send,
            release=self.bot.version,
            integrations=[SqlalchemyIntegration()],
        )

    @command(name="testsentry", permission="admin", show_in_help_list=False)
    def testsentry(self, mask, target, args):
        """Force an exception to test Sentry.

        %%testsentry
        """
        raise Exception("Sentry test")


import os
import pulumi
import pulumi_aws as aws
import json
from pulumi_terraform import state

config = pulumi.Config()
statefile = config.require("statefile")
script_dir = os.path.dirname(os.path.realpath(__file__))

# read outputs from Terraform
s = state.RemoteStateReference(
    "localstate",
    "local",
    state.LocalBackendArgs(path=os.path.join(f"{script_dir}/../tf", statefile)),
)

terraform_owned = aws.ec2.Vpc(
    "terraform_owned",
    cidr_block="10.0.0.0/16",
    enable_dns_hostnames=True,
    tags={
        "Name": "lbriggs",
    },
    opts=pulumi.ResourceOptions(protect=True),
)

# return the outputs from the terraform state
vpc_id = s.get_output("vpc_id")
private_subnet_ids = s.get_output("private_subnets")
public_subnet_ids = s.get_output("public_subnets")

pulumi.export("vpcId", vpc_id)
pulumi.export("publicSubnetIds", private_subnet_ids)
pulumi.export("privateSubnetIds", public_subnet_ids)

#
# Time to use Pulumi!
#

# Create an ECS cluster to run a container-based service.
cluster = aws.ecs.Cluster("cluster")

# create a security group referencing the VPC
group = aws.ec2.SecurityGroup(
    "web-secgrp",
    vpc_id=vpc_id,
    description="Enable HTTP access",
    ingress=[
        aws.ec2.SecurityGroupIngressArgs(
            protocol="tcp",
            from_port=80,
            to_port=80,
            cidr_blocks=["0.0.0.0/0"],
        )
    ],
    egress=[
        aws.ec2.SecurityGroupEgressArgs(
            protocol="-1",
            from_port=0,
            to_port=0,
            cidr_blocks=["0.0.0.0/0"],
        )
    ],
)

# Create a load balancer to listen for HTTP traffic on port 80.
# Reads the subnet IDs from Terraform state
alb = aws.lb.LoadBalancer(
    "app-lb",
    security_groups=[group.id],
    subnets=public_subnet_ids,
)

# create a target group
atg = aws.lb.TargetGroup(
    "app-tg",
    port=80,
    protocol="HTTP",
    target_type="ip",
    vpc_id=vpc_id,
)

# create a listener
wl = aws.lb.Listener(
    "web",
    load_balancer_arn=alb.arn,
    port=80,
    default_actions=[
        aws.lb.ListenerDefaultActionArgs(
            type="forward",
            target_group_arn=atg.arn,
        )
    ],
)

assume_policy_document = aws.iam.get_policy_document(
    statements=[
        aws.iam.GetPolicyDocumentStatementArgs(
            sid="AssumeRolePolicy",
            effect="Allow",
            principals=[
                aws.iam.GetPolicyDocumentStatementPrincipalArgs(
                    type="Service",
                    identifiers=["ecs-tasks.amazonaws.com"],
                )
            ],
            actions=["sts:AssumeRole"],
        )
    ]
)

# Create an IAM role that can be used by our service's task.
role = aws.iam.Role(
    "task-exec-role",
    assume_role_policy=assume_policy_document.json,
)

rpa = aws.iam.RolePolicyAttachment(
    "task-exec-policy",
    role=role.name,
    policy_arn="arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy",
)

# Spin up a load balanced service running our container image.
task_definition = aws.ecs.TaskDefinition(
    "app-task",
    family="terraform-state-demo",
    cpu="256",
    memory="512",
    network_mode="awsvpc",
    requires_compatibilities=["FARGATE"],
    execution_role_arn=role.arn,
    container_definitions=json.dumps(
        [
            {
                "name": "nginx",
                "image": "nginx",
                "portMappings": [
                    {"containerPort": 80, "hostPort": 80, "protocol": "tcp"}
                ],
            }
        ]
    ),
)

service = aws.ecs.Service(
    "app-svc",
    cluster=cluster.arn,
    desired_count=3,
    launch_type="FARGATE",
    task_definition=task_definition.arn,
    network_configuration=aws.ecs.ServiceNetworkConfigurationArgs(
        assign_public_ip=True,
        subnets=private_subnet_ids,
        security_groups=[group.id],
    ),
    load_balancers=[
        aws.ecs.ServiceLoadBalancerArgs(
            target_group_arn=atg.arn,
            container_name="nginx",
            container_port=80,
        )
    ],
    opts=pulumi.ResourceOptions(depends_on=[wl]),
)
pulumi.export("url", alb.dns_name)


import data_14dec2018
from math import floor

data = data_14dec2018.data

# Given an integer returns a list containing its digits in the same order
def int2digits(n):
    if n == 0:
        return [0]

    digits = []
    while n > 0:
        digits.insert(0, n % 10)
        n = floor(n / 10)

    return digits

# Given the current recipes returns the updated scores and the new recipes
def combine_recipes(scores, current_recipes, current_length):
    # Compute score associated with the current recipes
    sum_recipes_scores = sum([scores[recipe] for recipe in current_recipes])

    # Add newly created recipes
    append_scores = int2digits(sum_recipes_scores)
    start = current_length
    end = current_length + len(append_scores)
    new_scores = scores
    new_scores[start:end] = append_scores
    new_current_length = current_length + len(append_scores)

    # Compute new elves recipes
    new_current_recipes = []
    for recipe in current_recipes:
        new_recipe = (recipe + new_scores[recipe] + 1) % new_current_length
        new_current_recipes.append(new_recipe)

    return new_scores, new_current_recipes, new_current_length

# Number of elves
n_elves = 2

# Output info
start_point = data
length = 10
end_point = start_point + length

# Initial recipe for each elf
current_recipes = list(range(0, n_elves))

# Initial scores from data
scores = int2digits(37) + ([None] * (start_point + length))
current_length = 2

##############
### Part 1 ###
##############

# Run game for n_turns
while current_length < end_point:
    # Combine current recipes
    scores, current_recipes, current_length = combine_recipes(scores,
                                                              current_recipes,
                                                              current_length)

# Print output for part 1
print('Scores = ' + ''.join([str(d) for d in scores[start_point:end_point]]))

##############
### Part 2 ###
##############

# String to match
match = data
match = int2digits(match)
nn = len(match)

# Re-initialize
preallocate = int(2.5e7)
current_recipes = list(range(0, n_elves))
current_length = 2
scores = [None] * (preallocate + current_length)
scores[0:current_length] = int2digits(37)

# Run game until match appears
while current_length < preallocate:
    # Combine current recipes
    scores, current_recipes, current_length = combine_recipes(scores,
                                                              current_recipes,
                                                              current_length)
    if current_length >= nn:
        if scores[current_length-nn:current_length] == match:
            found = True
            print('Matched after ' + str(current_length-nn) + ' recipes')
            break
        elif scores[current_length-nn-1:current_length-1] == match:
            found = True
            print('Matched after ' + str(current_length-nn-1) + ' recipes')
            break
        else:
            found = False
            continue

if not found:
    print('Not found...')


#!/usr/bin/env python
"""
Simple program to verify we can read from a USB camera.
"""
import argparse
import cv2
from context import lib  # flake8: noqa
from lib.util import is_pressing

def run(channel):
    """
    Start the video capture on the given 'channel' (port). The number we pass
    is the camera number starting with 0 (typically for a built-in camera).
    External cameras are often on 1.
    """
    cap = cv2.VideoCapture(channel)

    while True:
        _, frame = cap.read()  # Capture frame-by-frame

        # Our operations on the frame come here
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Display the resulting frame
        cv2.imshow('frame', gray)
        if is_pressing('q'):
            break

    # When everything done, release the capture
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    PARSER = argparse.ArgumentParser(description=__doc__)
    PARSER.add_argument('-p', '--channel', default=1, type=int,
                        help='the USB channel containing camera, 0, 1, or 2')
    ARGS = PARSER.parse_args()

    print("""
    This program reads image frames from a connected camera, converts them to
    gray-scale, and displays them in a window. Call this program with a
    `--channel` option with a number. The number we pass is the camera number
    starting with 0 (typically for a built-in camera). External cameras are
    often on 1.

    Press 'q' to quit and close the window.
    """)

    run(ARGS.channel)


def fun1(maker,model,**others):
        others['maker'] = maker
        others['model'] = model
        for key,value in others.items():
                others[key] = value
        return others

car1 = fun1('BENZ',"S600",price = '3,000,000',length = '5m')
print(car1)


class Solution:
    def rearrangeCharacters(self, s: str, target: str) -> int:
        hashmap = dict()
        a = []
        for i in s:
            if i not in hashmap:
                hashmap[i] = 1
            else:
                hashmap[i] += 1

        for i in target:
            a.append(hashmap[i])

        return min(a) // len(target)

s = Solution()
print(s.rearrangeCharacters("abbaccaddaeea", "aaaaa"))


from typing import cast

from meiga import (
    Error,
    Result,
)
from PyQt5.QtGui import (
    QIcon,
)
from PyQt5.QtWidgets import (
    QMessageBox,
    QWidget,
)

from jenot import (
    logo,
    PollResult,
)
from jenot.args import Args
import jenot.notify as jnotify
from jenot.qt.processor import Processor

class MainWidget(QWidget):
    args: Args
    processors: list[Processor] = []

    def __init__(self) -> None:
        super().__init__()
        try:
            self.args = Args.parse(build_required=False)
        except SystemExit:
            QMessageBox.critical(None, 'Startup failed',
                "Could not parse some of required parameters, check your .jenotrc")
            raise
        self.jenot_icon = QIcon(logo('png'))

    def notify(self, result: Result[PollResult, Error], url: str) -> None:
        # have to store msgbox somewhere so python won't dispose of it
        self.msgbox = jnotify.qt(url, result)
        self.msgbox.show()

        if self.args.telegram:
            jnotify.telegram(url, result)

    def register(self, p: Processor) -> None:
        self.processors.append(p)
        p.finished.connect(self.notify)
        p.finished.connect(lambda: self.unregister(p))
        p.start()

    def unregister(self, p: Processor) -> None:
        self.processors.remove(p)


def primes():
    D = {}
    num = 2
    while num <= 1000:
        if num not in D:
            yield num
            D[num*num] = [num]
        else:
            for p in D[num]:
                D.setdefault(p + num, []).append(p)
            del D[num]
        num += 1

if __name__ == "__main__":
    for i in primes():
        print(i)
        if(i > 19):
            break