from dcore._dispatcher import *
from itertools import *
import numpy
import os
import copy

from ..tools import *

class SolverBase(object):
    """
    This class define the common interface of all solvers.
    All solvers must inherit from this class!
    """

    def __init__(self, beta, gf_struct, u_mat, n_iw=1025):
        """
        Initialize the solver.

        Parameters
        ----------
        beta : scalar
               Inverse temperature.
        gf_struct : dict{str:list}
                    Structure of the Green's functions. It must be a
                    dictionary, which maps the name of each block of the
                    Green's function as a string to a list of integer
                    indices.
                    For example: ``{'up': [1,2,3], 'down', [1,2,3]}``.
                    Allowed block names are 'up', 'down' and 'ud'.
        u_mat numpy.ndarray:
            four-index spin-full U-matrix.
            The index of each axis runs as (up, orb1), (up, orb2), ..., (down, orb1), (down, orb2), ...
        n_iw : integer, optional
               Number of Matsubara frequencies used for the Green's functions.
               Used for input and output.
        """

        assert isinstance(gf_struct, dict)

        if len(gf_struct) > 2:
            raise RuntimeError("gf_struct cannot contain more than two blocks.")

        allowed_block_names = ['up', 'down', 'ud']
        if not numpy.all([name in allowed_block_names for name in list(gf_struct.keys())]):
            raise RuntimeError("Allowed block names are up, down, ud.")

        # Member variables
        self.beta = beta
        self.gf_struct = gf_struct
        self.u_mat = u_mat
        self.n_iw = n_iw
        self.dims = numpy.array([len(indices) for name, indices in list(self.gf_struct.items())])
        self.n_flavors = numpy.sum(self.dims)
        self.use_spin_orbit = (len(gf_struct) == 1)
        self.n_orb = int(self.n_flavors/2)

        self.block_names = gf_block_names(self.use_spin_orbit) # This should be consistent with gf_struct.

        if self.u_mat.shape[0] != 2*self.n_orb:
            raise RuntimeError("Size of U matrix is wrong!")

        # Input data
        # self.G0_iw must be set to proper values before solve() is called.
        self._G0_iw = make_block_gf(GfImFreq, gf_struct, beta, n_iw)

        # Output data
        # These objects must be filled with results in solve().
        # Default value of Sigma_iw is 0, which will be used for the first iteration of DMFT self-consistent procedure.
        self._Sigma_iw = make_block_gf(GfImFreq, gf_struct, beta, n_iw)
        self._Sigma_iw.zero()
        self._Gimp_iw = make_block_gf(GfImFreq, gf_struct, beta, n_iw)
        self._Sigma_w = None
        #"self._G_l = None # Please define it if Legendre basis is used

        # Optional output data
        # This object will be saved into HDF5 file.
        self.quant_to_save = {}
        self.quant_to_save['nn_equal_time'] = numpy.zeros((self.n_flavors, self.n_flavors), dtype=float)

    def name(self):
        return "Base solver"

    def set_G0_iw(self, new_G0_iw):
        self._G0_iw << new_G0_iw.copy()

    def get_G0_iw(self):
        return self._G0_iw.copy()

    def get_Sigma_iw(self):
        return self._Sigma_iw.copy()

    def get_Gimp_iw(self):
        return self._Gimp_iw.copy()

    def get_Sigma_w(self):
        if self._Sigma_w is None:
            return None
        return self._Sigma_w.copy()

    def solve(self, rot, mpirun_command, params):
        """
        Solve the impudrity problem.

        self.G0_iw must be set before this function is called.

        INPUT:
          rot is a unitary matrix representing a basis transformation to reduce sign problem.
          If rot is None, no basis transformation is disabled.

          mpirun_command is a string (e.g. 'mpirun -np 8')

          params is a dict containing parameters.
          The following parameters must be set.
          Each impurity solver may require additional parameters.

              'random_seed_offset'  : int, off set value for random seed (mandatory)

              'calc_Sigma_w'        : bool, if real-frequency Sigma is computed (optional)
              'omega_min'           : float, lower limit of omega (must be set if calc_Sigma_iw is on).
              'omega_max'           : float, uppper limit of omega (must be set if calc_Sigma_iw is on).
              'n_omega'             : int, number of frequency points (must be set if calc_Sigma_iw is on).

        OUTPUT:
        ``Gimp_iw``, ``Sigma_iw`` must be calculated.

        """

        return NotImplementedError()

        # Set self.Gimp_iw, self.G_tau, self.Sigma_iw

    def calc_Xloc_ph(self, rot, mpirun_command, num_wf, num_wb, params_kw):
        """
        Compute Xloc(m, n, n') in p-h channel
                and chi_loc(m) (optional)

        Parameters
        ----------
        num_wf:
            Number of non-negative fermionic frequencies
        num_wb:
            Number of non-negative bosonic frequencies

        The other parameters are the same as for solve().

        Returns
        -------

        Xloc(m, n, n') : 3d ndarray of complex type
            Data for -num_wf <= n, n' < num_wf and m = 0, 1, ..., num_wb-1.

        chi_loc(m) : 1d ndarray of complex type
            return None if not computed

        """
        return NotImplementedError()

    def calc_Xloc_ph_sparse(self, rot, mpirun_command, wsample_ph, num_wb, params_kw):
        """
        Compute Xloc(m, n, n') in p-h channel only for specified frequency points
                and chi_loc(m) (optional)

        Parameters
        ----------
        rot
        mpirun_command:
            The same as solve()

        wsample_ph : 3*numpy.ndarray[N, 3]
           Sampling frequencies (in the order of fermion, fermion, boson)
           Fermion/boson frequencies are indexed by even/add integer numbers.

        num_wb:
            for chi_loc

        params_kw:
            The same as solve()

        Returns
        -------
        Xloc: ndarray or dict
           ndarray: (n_flavors, n_flavors, n_flavors, n_flavors, n_freqs)
           dict: key is (flavor0, flavor1, flavor2, flavor3), val is a ndarray of size (n_freqs,)
        chi_loc

        """
        return NotImplementedError()

    def calc_G2loc_ph_sparse(self, rot, mpirun_command, wsample_ph, params_kw):
        """
        Compute G2(m, n, n') in p-h channel only for specified frequency points.
        The definition of G2 is given in Eq. (3) of note/bse.pdf.

        Parameters
        ----------
        rot
        mpirun_command:
            The same as solve()

        wsample_ph : 3*numpy.ndarray[N, 3]
           Sampling frequencies (in the order of fermion, fermion, boson)
           Fermion/boson frequencies are indexed by even/add integer numbers.

        params_kw:
            The same as solve()

        Returns
        -------
        G2loc: ndarray or dict
           ndarray: (n_flavors, n_flavors, n_flavors, n_flavors, n_freqs)
           dict: key is (flavor0, flavor1, flavor2, flavor3), val is a ndarray of size (n_freqs,)

        """
        Xloc, _ =  self.calc_Xloc_ph_sparse(rot, mpirun_command, wsample_ph, 0, params_kw)
        # Convert Xloc to G2 using Eq. (4) in note/bse.pdf:
        #   G2loc_{pqrs} = beta * Xloc_{qprs}
        if isinstance(Xloc, numpy.ndarray):
            return self.beta * Xloc.transpose((1,0,2,3,4))
        elif isinstance(Xloc, dict):
            G2loc = {}
            for k, v in Xloc.items():
                G2loc[(k[1], k[0], k[2], k[3])] = self.beta * v
            return G2loc
        else:
            raise ValueError("Xloc must be a ndarray or a dict!")

    @classmethod
    def is_gf_realomega_available(cls):
        return False

    @classmethod
    def is_Floc_computable(cls):
        """ Can compute local full vertex """
        return False

def rotate_basis(rot, use_spin_orbit, u_matrix, Gfs=[], direction='forward'):
    """
    Rotate all Gf-like objects and U-matrix to the basis defined by rot

    :param direction:  'forward' or 'backward'
    :return u_matrix: rotated U matrix
    """

    if direction == 'forward':
        return _rotate_basis(rot, u_matrix, use_spin_orbit, Gfs)
    elif direction == 'backward':
        rot_conj_trans = {}
        for name, r in list(rot.items()):
            rot_conj_trans[name] = r.conjugate().transpose()
        return _rotate_basis(rot_conj_trans, u_matrix, use_spin_orbit, Gfs)
    else:
        raise RuntimeError("Unknown direction " + direction)

def _rotate_basis(rot, u_matrix, use_spin_orbit, Gfs):
    """
    Rotate all Gf-like object and U matrix to a new local basis defined by "rot".
    :param rot: matrix
    :return u_matrix: rotated U matrix
    """

    if use_spin_orbit:
        rot_spin_full = rot['ud']
    else:
        n_orb = rot['up'].shape[0]
        rot_spin_full = numpy.zeros((2*n_orb, 2*n_orb), dtype=complex)
        rot_spin_full[0:n_orb, 0:n_orb] = rot['up']
        rot_spin_full[n_orb:, n_orb:] = rot['down']

    for G in Gfs:
        for bname, gf in G:
            gf.from_L_G_R(rot[bname].transpose().conjugate(), gf, rot[bname])

    if not u_matrix is None:
        return numpy.einsum("ijkl,im,jn,ko,lp", u_matrix,
                                    numpy.conj(rot_spin_full), numpy.conj(rot_spin_full), rot_spin_full, rot_spin_full)

class PytriqsMPISolver(SolverBase):

    def __init__(self, beta, gf_struct, u_mat, n_iw=1025):
        """

        Initialize the solver.

        """

        super(PytriqsMPISolver, self).__init__(beta, gf_struct, u_mat, n_iw)

    def _impl_module_name(self):
        return ""

    def solve(self, rot, mpirun_command, params_kw):

        params = copy.deepcopy(params_kw)

        # Write input parameters
        with HDFArchive(os.path.abspath('input.h5'), 'w') as h:
            h['beta'] = self.beta
            h['gf_struct'] = self.gf_struct
            h['u_mat'] = self.u_mat
            h['n_iw'] = self.n_iw
            if not rot is None:
                h['rot'] = rot
            h['G0_iw'] = self._G0_iw
            h['params'] = params

            if 'calc_Sigma_w' in params and params['calc_Sigma_w']:
                h['calc_Sigma_w'] = True
                for k in ['omega_min', 'omega_max', 'n_omega']:
                    h[k] = params[k]
            else:
                h['calc_Sigma_w'] = False

        # Run a working horse
        commands = [sys.executable, "-m", self._impl_module_name()]
        commands.append(os.path.abspath('./input.h5'))
        commands.append(os.path.abspath('./output.h5'))

        with open('./output', 'w') as output_file:
            launch_mpi_subprocesses(mpirun_command, commands, output_file)
        with open('./output', 'r') as output_file:
            for line in output_file:
                print(line)

        # Read results
        with HDFArchive(os.path.abspath('output.h5'), 'r') as h:
            self._Sigma_iw << h['Sigma_iw']
            self._Gimp_iw << h['Gimp_iw']
            if 'Sigma_w' in h:
                self._Sigma_w = h['Sigma_w']

    def name(self):
        return "PytriqsMPISolver"

def compute_basis_rot(basis_rot, solver):
    """ Compute rotation matrix """
    nflavors = solver.n_flavors

    if basis_rot == 'None':
        rot = None
    elif basis_rot == 'Hloc':
        rot = compute_diag_basis(solver.get_G0_iw())
    else:
        if not os.path.exists(basis_rot):
            raise RuntimeError("Invalid basis_rot : {}".format(basis_rot))
        if solver.use_spin_orbit:
            rot = numpy.zeros((1, nflavors, nflavors), dtype=numpy.complex)
            read_potential(basis_rot, rot)
            rot = {'ud' : rot[0,:,:]}
        else:
            rot = numpy.zeros((2, nflavors//2, nflavors//2), dtype=numpy.complex)
            read_potential(basis_rot, rot)
            rot = {'up' : rot[0,:,:], 'down': rot[1,:,:]}
    return rot


#compute a map of the standard deviation of an array of pixels

from ser import *
import sys
from astropy.io import fits

fid = Ser(sys.argv[1])

mean = fid.load_img(0) * 1.0
N = fid.count
print(N)
print("count ", fid.count)
for idx in range(1, N):
    sum = fid.load_img(idx)
    mean = mean + sum
    print(idx, np.mean(sum), sum)
    if (idx % 100 == 0):
        print(idx)
fid.close()

mean = mean / N
hdr = fits.header.Header()
fits.writeto(sys.argv[2], (mean).astype(np.float32), hdr, overwrite=True)


peca1 = str(input())
peca2 = str(input())
# Giving the variables
valoresPeca1 = peca1.split(" ")
valoresPeca2 = peca2.split(" ")
# Reutilizing him
valoresPeca1 = float(valoresPeca1[1])*float(valoresPeca1[2])
valoresPeca2 = float(valoresPeca2[1])*float(valoresPeca2[2])
# display the result
print("VALOR A PAGAR: "
      f"R$ {valoresPeca1 + valoresPeca2:.2f}")


import os
from tkinter import messagebox
import sys

en="/home/"+sys.argv[1]+"/.by-mukonqi/metterxp/language/en.txt"
tr="/home/"+sys.argv[1]+"/.by-mukonqi/metterxp/language/tr.txt"
debian="/etc/debian_version"
fedora="/etc/fedora-release"
solus="/etc/solus-release"

if not os.path.isdir("/home/"+str(sys.argv[1])+"/.by-mukonqi/metterxp/"):
    messagebox.showerror("User Error - Fatal | MetterXP","Please start only MetterXP (just select MetterXP in the applications section or type metterxp) to set up language and theme preferences.")
    exit()

if not os.getuid() == 0:
    if os.path.isfile(en):
        messagebox.showerror("Error","Only root can run this module!")
        exit()
    elif os.path.isfile(tr):
        messagebox.showerror("Hata","Sadece kök kullanıcı bu modülü çalıştırabilir!")
        exit()

if os.path.isfile(debian):
    os.system("apt upgrade -y")
elif os.path.isfile(fedora):
    os.system("dnf update -y")
elif os.path.isfile(solus):
    os.system("eopkg up -y")
if os.path.isfile(en):
    messagebox.showinfo("Information","If a system update is available, it has been completed.")
elif os.path.isfile(tr):
    messagebox.showinfo("Bilgilendirme","Eğer bir sistem güncellemesi mevcutsa tamamlandı.")
exit()


from django.shortcuts import render, HttpResponse
from django.http import HttpResponse
import os

import pytube
import ssl
from django.http import FileResponse
# Create your views here.
def home(request):
    if request.method =='POST':
        # url = request.POST['q']

        result = download_video('https://www.youtube.com/watch?v=5YZZSsYSp88')

        print(f"download {result}")

    return render(request, 'yt_downloader/home.html')

def download_video(url):

    if url:
        try:

            ssl._create_default_https_context = ssl._create_unverified_context

            # youtube = pytube.YouTube(url)
            # video = pytube.YouTube(url).streams.get_highest_resolution()
            video = pytube.YouTube(url).streams.get_lowest_resolution()

            # video = youtube.streams.first()
            out_file = video.download()

            base, ext = os.path.splitext(out_file)
            new_file = base + '.mp3'
            os.rename(out_file, new_file)

            print("out_file:" , out_file)
            print('*' * 20)
            print("new file:", new_file)
            print("*" * 20)
            print('video:', video)

            try:
                # with open('my_video.mp4', 'wb') as f:
                #     f.write(new_file)
                rea_response = HttpResponse(new_file, content_type='video/mp4')
                rea_response['Content-Disposition'] = 'attachment; filename=my_video.mp4'
                return rea_response
            except TimeoutError:
                return HttpResponse(TimeoutError)
            # print(video)
            #
            # return video

        except ValueError:
            print(ValueError.args)
            print('Error')

            return False


import gobject
import gtk
import re
from shell_buffer import ShellBuffer, ADJUST_NONE, ADJUST_BEFORE, ADJUST_AFTER
from chunks import StatementChunk, CommentChunk, BlankChunk
from completion_popup import CompletionPopup
from doc_popup import DocPopup
from notebook import NotebookFile
import sanitize_textview_ipc

class ShellView(gtk.TextView):
    __gsignals__ = {
        'backspace' : 'override',
        'expose-event': 'override',
        'focus-out-event': 'override',
        'button-press-event': 'override',
        'button-release-event': 'override',
        'motion-notify-event': 'override',
        'key-press-event': 'override',
        'leave-notify-event': 'override',
        'motion-notify-event': 'override',
        'realize': 'override',
        'unrealize': 'override',
        'size-allocate': 'override'
   }

    def __init__(self, buf):
        buf.worksheet.connect('chunk-inserted', self.on_chunk_inserted)
        buf.worksheet.connect('chunk-changed', self.on_chunk_changed)
        buf.worksheet.connect('chunk-status-changed', self.on_chunk_status_changed)
        buf.worksheet.connect('notify::state', self.on_notify_state)
        buf.connect('add-custom-result', self.on_add_custom_result)
        buf.connect('pair-location-changed', self.on_pair_location_changed)

        gtk.TextView.__init__(self, buf)
        self.set_border_window_size(gtk.TEXT_WINDOW_LEFT, 10)
        self.set_left_margin(2)

        # Attach a "behavior object" to the view which, by ugly hacks, makes it
        # do simply and reasonable things for cut-and-paste and DND
        sanitize_textview_ipc.sanitize_view(self)

        self.add_events(gtk.gdk.LEAVE_NOTIFY_MASK)

        self.__completion_popup = CompletionPopup(self)
        self.__doc_popup = DocPopup()
        self.__mouse_over_object = None
        self.__mouse_over_timeout = None

        buf = self.get_buffer()
        self.__mouse_over_start = buf.create_mark(None, buf.get_start_iter(), True)

    def __get_worksheet_line_yrange(self, line):
        buffer_line = self.get_buffer().pos_to_iter(line)
        return self.get_line_yrange(buffer_line)

    def __get_worksheet_line_at_y(self, y, adjust):
        buf = self.get_buffer()
        (buffer_line, _) = self.get_line_at_y(y)
        return buf.iter_to_pos(buffer_line, adjust)[0]

    def paint_chunk(self, cr, area, chunk, fill_color, outline_color):
        buf = self.get_buffer()

        (y, _) = self.__get_worksheet_line_yrange(chunk.start)
        (end_y, end_height) = self.__get_worksheet_line_yrange(chunk.end - 1)
        height = end_y + end_height - y

        (_, window_y) = self.buffer_to_window_coords(gtk.TEXT_WINDOW_LEFT, 0, y)
        cr.rectangle(area.x, window_y, area.width, height)
        cr.set_source_rgb(*fill_color)
        cr.fill()

        cr.rectangle(0.5, window_y + 0.5, 10 - 1, height - 1)
        cr.set_source_rgb(*outline_color)
        cr.set_line_width(1)
        cr.stroke()

    def do_realize(self):
        gtk.TextView.do_realize(self)

        self.get_window(gtk.TEXT_WINDOW_LEFT).set_background(self.style.white)

        # While the the worksheet is executing, we want to display a watch cursor
        # Trying to override the cursor setting of GtkTextView is really hard because
        # of things like hiding the cursor when typing, so we take the simple approach
        # of using an input-only "cover window" that we set the cursor on and
        # show on top of the GtkTextView's normal window.

        self.__watch_window = gtk.gdk.Window(self.window,
                                             self.allocation.width, self.allocation.height,
                                             gtk.gdk.WINDOW_CHILD,
                                             (gtk.gdk.SCROLL_MASK |
                                              gtk.gdk.BUTTON_PRESS_MASK |
                                              gtk.gdk.BUTTON_RELEASE_MASK |
                                              gtk.gdk.POINTER_MOTION_MASK |
                                              gtk.gdk.POINTER_MOTION_HINT_MASK),
                                             gtk.gdk.INPUT_ONLY,
                                             x=0, y=0)
        self.__watch_window.set_cursor(gtk.gdk.Cursor(gtk.gdk.WATCH))
        self.__watch_window.set_user_data(self)

        if self.get_buffer().worksheet.state == NotebookFile.EXECUTING:
            self.__watch_window.show()
            self.__watch_window.raise_()

    def do_unrealize(self):
        self.__watch_window.set_user_data(None)
        self.__watch_window.destroy()
        self.__watch_window = None
        gtk.TextView.do_unrealize(self)

    def do_size_allocate(self, allocation):
        gtk.TextView.do_size_allocate(self, allocation)
        if (self.flags() & gtk.REALIZED) != 0:
            self.__watch_window.resize(allocation.width, allocation.height)

    def __expose_window_left(self, event):
        (_, start_y) = self.window_to_buffer_coords(gtk.TEXT_WINDOW_LEFT, 0, event.area.y)
        start_line = self.__get_worksheet_line_at_y(start_y, adjust=ADJUST_AFTER)

        (_, end_y) = self.window_to_buffer_coords(gtk.TEXT_WINDOW_LEFT, 0, event.area.y + event.area.height - 1)
        end_line = self.__get_worksheet_line_at_y(end_y, adjust=ADJUST_BEFORE)

        buf = self.get_buffer()

        cr = event.window.cairo_create()

        for chunk in buf.worksheet.iterate_chunks(start_line, end_line + 1):
            if isinstance(chunk, StatementChunk):
                if chunk.error_message != None:
                    self.paint_chunk(cr, event.area, chunk, (1, 0, 0), (0.5, 0, 0))
                elif chunk.needs_compile:
                    self.paint_chunk(cr, event.area, chunk, (1, 1, 0), (0.5, 0.5, 0))
                elif chunk.needs_execute:
                    self.paint_chunk(cr, event.area, chunk, (1, 0, 1), (0.5, 0.5, 0))
                else:
                    self.paint_chunk(cr, event.area, chunk, (0, 0, 1), (0, 0, 0.5))

    def __expose_pair_location(self, event):
        pair_location = self.get_buffer().get_pair_location()
        if pair_location == None:
            return

        rect = self.get_iter_location(pair_location)

        rect.x, rect.y = self.buffer_to_window_coords(gtk.TEXT_WINDOW_TEXT, rect.x, rect.y)

        if (rect.y + rect.height <= event.area.y or rect.y >= event.area.y + event.area.height):
            return

        cr = event.window.cairo_create()
        cr.set_line_width(1.)
        cr.rectangle(rect.x + 0.5, rect.y + 0.5, rect.width - 1, rect.height - 1)
        cr.set_source_rgb(0.6, 0.6, 0.6)
        cr.stroke()

    def do_expose_event(self, event):
        if event.window == self.get_window(gtk.TEXT_WINDOW_LEFT):
            self.__expose_window_left(event)
            return False

        gtk.TextView.do_expose_event(self, event)

        if event.window == self.get_window(gtk.TEXT_WINDOW_TEXT):
            self.__expose_pair_location(event)

        return False

    def __get_line_text(self, iter):
        start = iter.copy()
        start.set_line_index(0)
        end = iter.copy()
        end.forward_to_line_end()

        return start.get_slice(end)

    # This is likely overengineered, since we're going to try as hard as possible not to
    # have tabs in our worksheets
    def __count_indent(self, text):
        indent = 0
        for c in text:
            if c == ' ':
                indent += 1
            elif c == '\t':
                indent += 8 - (indent % 8)
            else:
                break

        return indent

    def __find_outdent(self, iter):
        buf = self.get_buffer()
        line, _ = buf.iter_to_pos(iter)

        current_indent = self.__count_indent(buf.worksheet.get_line(line))

        while line > 0:
            line -= 1
            line_text = buf.worksheet.get_line(line)

            indent = self.__count_indent(line_text)
            if indent < current_indent:
                return re.match(r"^[\t ]*", line_text).group(0)

        return ""

    def __find_default_indent(self, iter):
        buf = self.get_buffer()
        line, offset = buf.iter_to_pos(iter)

        while line > 0:
            line -= 1
            chunk = buf.worksheet.get_chunk(line)
            if isinstance(chunk, StatementChunk):
                return chunk.tokenized.get_next_line_indent(line - chunk.start)
            elif isinstance(chunk, CommentChunk):
                return " " * self.__count_indent(buf.worksheet.get_line(line))

        return ""

    def __reindent_line(self, iter, indent_text):
        buf = self.get_buffer()

        line_text = self.__get_line_text(iter)
        prefix = re.match(r"^[\t ]*", line_text).group(0)

        diff = self.__count_indent(indent_text) - self.__count_indent(prefix)
        if diff == 0:
            return 0

        common_len = 0
        for a, b in zip(prefix, indent_text):
            if a != b:
                break
            common_len += 1

        start = iter.copy()
        start.set_line_offset(common_len)
        end = iter.copy()
        end.set_line_offset(len(prefix))

        # Nitpicky-detail. If the selection starts at the start of the line, and we are
        # inserting white-space there, then the whitespace should be *inside* the selection
        mark_to_start = None
        if common_len == 0 and buf.get_has_selection():
            mark = buf.get_insert()
            if buf.get_iter_at_mark(mark).compare(start) == 0:
                mark_to_start = mark

            mark = buf.get_selection_bound()
            if buf.get_iter_at_mark(mark).compare(start) == 0:
                mark_to_start = mark

        buf.delete(start, end)
        buf.insert(end, indent_text[common_len:])

        if mark_to_start != None:
            end.set_line_offset(0)
            buf.move_mark(mark_to_start, end)

        return diff

    def __reindent_selection(self, outdent):
        buf = self.get_buffer()

        bounds = buf.get_selection_bounds()
        if bounds == ():
            insert_mark = buf.get_insert()
            bounds = buf.get_iter_at_mark(insert_mark), buf.get_iter_at_mark(insert_mark)
        start, end = bounds

        line, _ = buf.iter_to_pos(start, adjust=ADJUST_AFTER)
        end_line, end_offset = buf.iter_to_pos(end, adjust=ADJUST_BEFORE)
        if end_offset == 0 and end_line > line:
            end_line -= 1

        iter = buf.pos_to_iter(line)

        if outdent:
            indent_text = self.__find_outdent(iter)
        else:
            indent_text = self.__find_default_indent(iter)

        diff = self.__reindent_line(iter, indent_text)
        while True:
            line += 1
            if line > end_line:
                return

            iter = buf.pos_to_iter(line)
            current_indent = self.__count_indent(self.__get_line_text(iter))
            self.__reindent_line(iter, max(0, " " * (current_indent + diff)))

    def __hide_completion(self):
        if self.__completion_popup.showing:
            self.__completion_popup.popdown()

    def __update_completion(self):
        if self.__completion_popup.showing:
            self.__completion_popup.update()

    def do_focus_out_event(self, event):
        self.__hide_completion()
        self.__doc_popup.popdown()
        return gtk.TextView.do_focus_out_event(self, event)

    def __rewrite_window(self, event):
        # Mouse events on the "watch window" that covers the text view
        # during calculation need to be forwarded to the real text window
        # since it looks bad if keynav works, but you can't click on the
        # window to set the cursor, select text, and so forth

        if event.window == self.__watch_window:
            event.window = self.get_window(gtk.TEXT_WINDOW_TEXT)

    def do_button_press_event(self, event):
        self.__rewrite_window(event)

        self.__doc_popup.popdown()

        return gtk.TextView.do_button_press_event(self, event)

    def do_button_release_event(self, event):
        self.__rewrite_window(event)

        return gtk.TextView.do_button_release_event(self, event)

    def do_motion_notify_event(self, event):
        self.__rewrite_window(event)

        return gtk.TextView.do_motion_notify_event(self, event)

    def do_key_press_event(self, event):
        buf = self.get_buffer()

        if self.__completion_popup.focused and self.__completion_popup.on_key_press_event(event):
            return True

        if self.__doc_popup.focused:
            self.__doc_popup.on_key_press_event(event)
            return True

        if event.keyval in (gtk.keysyms.F2, gtk.keysyms.KP_F2):
            self.__hide_completion()

            if self.__doc_popup.showing:
                self.__doc_popup.focus()
            else:
                insert = buf.get_iter_at_mark(buf.get_insert())
                line, offset = buf.iter_to_pos(insert, adjust=ADJUST_NONE)
                if line != None:
                    obj, start_line, start_offset, _, _ = buf.worksheet.get_object_at_location(line, offset, include_adjacent=True)
                else:
                    obj = None

                if obj != None:
                    start = buf.pos_to_iter(start_line, start_offset)
                    self.__stop_mouse_over()
                    self.__doc_popup.set_target(obj)
                    self.__doc_popup.position_at_location(self, start)
                    self.__doc_popup.popup_focused()

            return True

        self.__doc_popup.popdown()

        if event.keyval in (gtk.keysyms.KP_Enter, gtk.keysyms.Return):
            self.__hide_completion()

            increase_indent = False
            insert = buf.get_iter_at_mark(buf.get_insert())
            line, pos = buf.iter_to_pos(insert, adjust=ADJUST_NONE)

            # Inserting return inside a ResultChunk would normally do nothing
            # but we want to make it insert a line after the chunk
            if line == None and not buf.get_has_selection():
                line, pos = buf.iter_to_pos(insert, adjust=ADJUST_BEFORE)
                iter = buf.pos_to_iter(line, -1)
                buf.place_cursor(iter)
                buf.insert_interactive(iter, "\n", True)

                return True

            buf.begin_user_action()

            gtk.TextView.do_key_press_event(self, event)
            # We need the chunks to be updated when computing the line indents
            buf.worksheet.rescan()

            insert = buf.get_iter_at_mark(buf.get_insert())

            self.__reindent_line(insert, self.__find_default_indent(insert))

            # If we have two comment lines in a row, assume a block comment
            if (line > 0 and
                isinstance(buf.worksheet.get_chunk(line), CommentChunk) and
                isinstance(buf.worksheet.get_chunk(line - 1), CommentChunk)):
                self.get_buffer().insert_interactive_at_cursor("# ", True)

            buf.end_user_action()

            return True
        elif event.keyval in (gtk.keysyms.Tab, gtk.keysyms.KP_Tab) and event.state & gtk.gdk.CONTROL_MASK == 0:
            buf.begin_user_action()
            self.__reindent_selection(outdent=False)
            buf.end_user_action()

            self.__update_completion()
            return True
        elif event.keyval == gtk.keysyms.ISO_Left_Tab and event.state & gtk.gdk.CONTROL_MASK == 0:
            buf.begin_user_action()
            self.__reindent_selection(outdent=True)
            buf.end_user_action()

            self.__update_completion()
            return True
        elif event.keyval == gtk.keysyms.space and event.state & gtk.gdk.CONTROL_MASK != 0:
            if self.__completion_popup.showing:
                self.__completion_popup.popdown()
            else:
                if self.__doc_popup.showing:
                    self.__doc_popup.popdown()
                self.__completion_popup.popup()
            return True
        elif event.keyval in (gtk.keysyms.z, gtk.keysyms.Z) and \
                (event.state & gtk.gdk.CONTROL_MASK) != 0 and \
                (event.state & gtk.gdk.SHIFT_MASK) == 0:
            buf.worksheet.undo()

            self.__update_completion()
            return True
        # This is the gedit/gtksourceview binding to cause your hands to fall off
        elif event.keyval in (gtk.keysyms.z, gtk.keysyms.Z) and \
                (event.state & gtk.gdk.CONTROL_MASK) != 0 and \
                (event.state & gtk.gdk.SHIFT_MASK) != 0:
            buf.worksheet.redo()

            self.__update_completion()
            return True
        # This is the familiar binding (Eclipse, etc). Firefox supports both
        elif event.keyval in (gtk.keysyms.y, gtk.keysyms.Y) and event.state & gtk.gdk.CONTROL_MASK != 0:
            buf.worksheet.redo()

            self.__update_completion()
            return True

        if gtk.TextView.do_key_press_event(self, event):
            self.__update_completion()
            return True
        else:
            return False

    def __show_mouse_over(self):
        self.__mouse_over_timeout = None

        if self.__completion_popup.showing:
            return

        self.__doc_popup.set_target(self.__mouse_over_object)
        location = self.get_buffer().get_iter_at_mark(self.__mouse_over_start)
        self.__doc_popup.position_at_location(self, location)
        self.__doc_popup.popup()

        return False

    def __stop_mouse_over(self):
        if self.__mouse_over_timeout:
            gobject.source_remove(self.__mouse_over_timeout)
            self.__mouse_over_timeout = None

        self.__mouse_over_object = None

    def do_motion_notify_event(self, event):
        if event.window == self.get_window(gtk.TEXT_WINDOW_TEXT) and not self.__doc_popup.focused:
            buf = self.get_buffer()

            iter, _ = self.get_iter_at_position(int(event.x), int(event.y))
            line, offset = buf.iter_to_pos(iter, adjust=ADJUST_NONE)
            if line != None:
                obj, start_line, start_offset, _,_ = buf.worksheet.get_object_at_location(line, offset)
            else:
                obj = None

            if not obj is self.__mouse_over_object:
                self.__stop_mouse_over()
                self.__doc_popup.popdown()
                if obj != None:
                    start = buf.pos_to_iter(start_line, start_offset)
                    buf.move_mark(self.__mouse_over_start, start)

                    self.__mouse_over_object = obj
                    try:
                        timeout = self.get_settings().get_property('gtk-tooltip-timeout')
                    except TypeError: # GTK+ < 2.12
                        timeout = 500
                    self.__mouse_over_timeout = gobject.timeout_add(timeout, self.__show_mouse_over)

        return gtk.TextView.do_motion_notify_event(self, event)

    def do_leave_notify_event(self, event):
        self.__stop_mouse_over()
        if not self.__doc_popup.focused:
            self.__doc_popup.popdown()
        return False

    def do_backspace(self):
        buf = self.get_buffer()

        insert = buf.get_iter_at_mark(buf.get_insert())
        line, offset = buf.iter_to_pos(insert)

        current_chunk = buf.worksheet.get_chunk(line)
        if isinstance(current_chunk, StatementChunk) or isinstance(current_chunk, BlankChunk):
            line_text = buf.worksheet.get_line(line)[0:offset]

            if re.match(r"^[\t ]+$", line_text):
                self.__reindent_line(insert, self.__find_outdent(insert))
                return

        return gtk.TextView.do_backspace(self)

    def __invalidate_status(self, chunk):
        buf = self.get_buffer()

        (start_y, start_height) = self.__get_worksheet_line_yrange(chunk.start)
        (end_y, end_height) = self.__get_worksheet_line_yrange(chunk.end - 1)

        (_, window_y) = self.buffer_to_window_coords(gtk.TEXT_WINDOW_LEFT, 0, start_y)

        if self.window:
            self.get_window(gtk.TEXT_WINDOW_LEFT).invalidate_rect((0, window_y, 10, end_y + end_height - start_y), False)

    def on_chunk_inserted(self, worksheet, chunk):
        self.__invalidate_status(chunk)

    def on_chunk_changed(self, worksheet, chunk, changed_lines):
        self.__invalidate_status(chunk)

    def on_chunk_status_changed(self, worksheet, chunk):
        self.__invalidate_status(chunk)

    def on_notify_state(self, worksheet, param_spec):
        if (self.flags() & gtk.REALIZED) != 0:
            if worksheet.state == NotebookFile.EXECUTING:
                self.__watch_window.show()
                self.__watch_window.raise_()
            else:
                self.__watch_window.hide()

    def on_add_custom_result(self, buf, result, anchor):
        widget = result.create_widget()
        widget.show()
        self.add_child_at_anchor(widget, anchor)

    def __invalidate_char_position(self, iter):
        y, height = self.get_line_yrange(iter)
        if self.window:
            text_window = self.get_window(gtk.TEXT_WINDOW_TEXT)
            width, _ = text_window.get_size()
            text_window.invalidate_rect((0, y, width, height), False)

    def on_pair_location_changed(self, buf, old_position, new_position):
        if old_position:
            self.__invalidate_char_position(old_position)
        if new_position:
            self.__invalidate_char_position(new_position)

    def calculate(self):
        buf = self.get_buffer()

        buf.worksheet.calculate()

        # This is a hack to work around the fact that scroll_mark_onscreen()
        # doesn't wait for a size-allocate cycle, so doesn't properly handle
        # embedded request widgets
        self.size_request()
        self.size_allocate((self.allocation.x, self.allocation.y,
                            self.allocation.width, self.allocation.height))

        self.scroll_mark_onscreen(buf.get_insert())

    def copy_as_doctests(self):
        buf = self.get_buffer()

        bounds = buf.get_selection_bounds()
        if bounds == ():
            start, end = buf.get_iter_at_mark(buf.get_insert())
        else:
            start, end = bounds

        start_line, start_offset = buf.iter_to_pos(start, adjust=ADJUST_BEFORE)
        end_line, end_offset = buf.iter_to_pos(end, adjust=ADJUST_BEFORE)

        doctests = buf.worksheet.get_doctests(start_line, end_line + 1)
        self.get_clipboard(gtk.gdk.SELECTION_CLIPBOARD).set_text(doctests)


from selenium import webdriver
from selenium.webdriver.common.by import By
import time

# link = "http://suninjuly.github.io/registration1.html" # old registration1.html
link = "http://suninjuly.github.io/registration2.html" #new registration2.html

try:
    browser = webdriver.Chrome()
    browser.get(link)

    input1 = browser.find_element(By.CSS_SELECTOR, ".first_block .form-control.first") # First name
    input1.send_keys("Ivan")
    input2 = browser.find_element(By.CSS_SELECTOR, ".first_block .form-control.second") # Last name
    input2.send_keys("Petrov")
    input3 = browser.find_element(By.CSS_SELECTOR, ".first_block .form-control.third") # Email
    input3.send_keys("ivan@petrov.com")

    # input4 = browser.find_element(By.CSS_SELECTOR, ".second_block .form-control.first") # Phone
    # input4.send_keys("1234567")
    #
    # input5 = browser.find_element(By.CSS_SELECTOR, ".second_block .form-control.second") # Address
    # input5.send_keys("Russia")

    button = browser.find_element(By.CLASS_NAME, "btn.btn-default") # Button
    button.click()

finally:
    time.sleep(5)
    browser.quit()


import numpy as np
import pandas as pd

from numpy import nan
from pandas import DataFrame, Series

def test_1(get_rating, ratings):
    print('Check output type: ', end='')
    _solution = get_rating(ratings, 182, 2710)
    _expected_type = float
    assert isinstance(_solution,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(_solution)}'
    print('success!')

    print('Check get_rating(): ', end='')
    for _user, _movie, _solution in [(182, 2571, 5.0), (182, 2710, 4.5),
                                     (182, 4306, 4.0), (195, 2571, 3.0),
                                     (195, 2710, 1.0), (195, 4306, 3.0),
                                     (204, 2571, 4.5), (204, 2710, 5.0),
                                     (204, 4306, 4.0), (376, 2571, 3.5),
                                     (376, 2710, 1.5), (376, 4306, 4.0),
                                     (542, 2571, 5.0), (542, 2710, 0.5),
                                     (542, 4306, 5.0), (182, 1234, np.nan),
                                     (195, 5678, np.nan)]:
        if np.isnan(_solution):
            assert np.isnan(get_rating(ratings, _user,
                                       _movie)), f'Return NaN if there is no rating for the user/movie combination'
        else:
            assert get_rating(ratings, _user,
                              _movie) == _solution, f'The rating for user {_user}, {_movie} should be {_solution}'
    print('Success!')

def test_2(pivot_ratings, ratings):
    print('Check output type: ', end='')
    _student_answer = pivot_ratings(ratings)
    _expected_type = pd.DataFrame
    assert isinstance(_student_answer,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(_solution)}'
    print('success!')

    print('Testing pivot on complete data set: ', end='')
    _solution = DataFrame([[5.0, 3.0, 4.5, 3.5, 5.0], [4.5, 1.0, 5.0, 1.5,
                                                       0.5],
                           [4.0, 3.0, 4.0, 4.0, 5.0]],
                          columns=[182, 195, 204, 376, 542],
                          index=[2571, 2710, 4306])
    _student_answer = pivot_ratings(ratings)

    pd.testing.assert_frame_equal(_student_answer, _solution,
                                  check_names=False, check_like=True)
    print('success!')

    print('Testing pivot on incomplete data set: ', end='')

    _incomplete_ratings = DataFrame(
        [[182, 2571, 5.0, 1054779786], [182, 2710, 4.5, 1063284735],
         [195, 4306, 3.0, 994032742], [204, 2571, 4.5, 1327183462],
         [204, 2710, 5.0, 1327185697], [376, 2571, 3.5, 1364994024],
         [376, 2710, 1.5, 1364994544], [542, 2571, 5.0, 1163386800],
         [542, 2710, 0.5, 1163387159]],
        columns=['userId', 'movieId', 'rating', 'timestamp'],
        index=[0, 1, 5, 6, 7, 9, 10, 12, 13])

    _solution = DataFrame(
        [[5.0, nan, 4.5, 3.5, 5.0], [4.5, nan, 5.0, 1.5, 0.5],
         [nan, 3.0, nan, nan, nan]],
        columns=[182, 195, 204, 376, 542],
        index=[2571, 2710, 4306])

    _student_answer = pivot_ratings(_incomplete_ratings)

    pd.testing.assert_frame_equal(_student_answer, _solution,
                                  check_names=False, check_like=True)
    print('success!')

def test_8(create_similarity_matrix_manhattan, utility_matrix):
    print('Check output type: ', end='')
    _student_answer = create_similarity_matrix_manhattan(utility_matrix)
    _expected_type = pd.DataFrame
    assert isinstance(_student_answer,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(_student_answer)}'
    print('success!')

    print('Testing create_similarity_matrix_manhattan on complete data set: ',
          end='')
    _student_answer = create_similarity_matrix_manhattan(utility_matrix)
    _solution = DataFrame([[1.0, 0.09523809523809523, 0.3333333333333333],
                           [0.09523809523809523, 1.0, 0.08695652173913043],
                           [0.3333333333333333, 0.08695652173913043, 1.0]],
                          columns=[2571, 2710, 4306], index=[2571, 2710, 4306])
    np.testing.assert_allclose(_student_answer, _solution, rtol=1e-1)
    print('success!')

    print(
        'Testing create_similarity_matrix_manhattan on incomplete data set: ',
        end='')
    _utility_matrix = DataFrame(
        [[5.0, nan, 4.5, nan, 5.0], [4.5, nan, 5.0, 1.5, 0.5],
         [nan, 3.0, nan, nan, nan]], columns=[182, 195, 204, 376, 542],
        index=[2571, 2710, 4306])
    _student_answer = create_similarity_matrix_manhattan(_utility_matrix)
    _solution = DataFrame(
        [[1.0, 0.15384615384615385, nan], [0.15384615384615385, 1.0, nan],
         [nan, nan, 1.0]], columns=[2571, 2710, 4306],
        index=[2571, 2710, 4306])

    np.testing.assert_allclose(_student_answer, _solution, rtol=1e-1)
    print('success!')

    print('\n\nOutput second test:')
    print(_student_answer)

def test_10(create_similarity_matrix_euclid, utility_matrix):
    print('Check output type: ', end='')
    _student_answer = create_similarity_matrix_euclid(utility_matrix)
    _expected_type = pd.DataFrame
    assert isinstance(_student_answer, _expected_type), f'expected output to be of type {_expected_type}, not {type(_expected_type)}'
    print('success!')

    print('Testing create_similarity_matrix_euclid on complete data set: ',
          end='')
    _student_answer = create_similarity_matrix_euclid(utility_matrix)
    _solution = DataFrame([[1.0, 0.1571856809867317, 0.4494897427831781],
                           [0.1571856809867317, 1.0, 0.15072240113145766],
                           [0.4494897427831781, 0.15072240113145766, 1.0]],
                          columns=[2571, 2710, 4306], index=[2571, 2710, 4306])
    np.testing.assert_allclose(_student_answer, _solution, rtol=1e-1)
    print('success!')

    print('Testing create_similarity_matrix_euclid on incomplete data set: ',
          end='')
    _utility_matrix = DataFrame(
        [[5.0, nan, 4.5, nan, 5.0], [4.5, nan, 5.0, 1.5, 0.5],
         [nan, 3.0, nan, nan, nan]], columns=[182, 195, 204, 376, 542],
        index=[2571, 2710, 4306])
    _student_answer = create_similarity_matrix_euclid(_utility_matrix)
    _solution = DataFrame(
        [[1.0, 0.18001097668719743, nan], [0.18001097668719743, 1.0, nan],
         [nan, nan, 1.0]], columns=[2571, 2710, 4306],
        index=[2571, 2710, 4306])
    np.testing.assert_allclose(_student_answer, _solution, rtol=1e-1)
    print('success!')

def test_11(create_similarity_matrix_cosine, utility_matrix):
    print('Check output type: ', end='')
    _student_answer = create_similarity_matrix_cosine(utility_matrix)
    _expected_type = pd.DataFrame
    assert isinstance(_student_answer, _expected_type), f'expected output to be of type {_expected_type}, not {type(_student_answer)}'
    print('success!')

    print('Testing create_similarity_matrix_cosine on complete data set: ',
          end='')
    _solution = DataFrame([[0.9999999999999999, 0.8347319075206693, 0.992843914076249],
                             [0.8347319075206693, 1.0, 0.7829084180122037],
                             [0.992843914076249, 0.7829084180122037, 0.9999999999999998]],
                            columns=[2571, 2710, 4306], index=[2571, 2710, 4306])
    _student_answer = create_similarity_matrix_cosine(utility_matrix)

    np.testing.assert_allclose(_student_answer, _solution, rtol=1e-1)
    print('success!')

    print('Testing create_similarity_matrix_cosine on incomplete data set: ',
          end='')

    _utility_matrix = DataFrame(
        [[5.0, nan, 4.5, nan, 5.0], [4.5, nan, 5.0, 1.5, 0.5],
         [nan, 3.0, nan, nan, nan]], columns=[182, 195, 204, 376, 542],
        index=[2571, 2710, 4306])

    _solution = DataFrame([[1.0000000000, 0.8347319075, 0.9928439141],
                           [0.8347319075, 1.0000000000, 0.7829084180],
                           [0.9928439141, 0.7829084180, 1.0000000000]],
                          columns=[2571, 2710, 4306], index=[2571, 2710, 4306])

    _student_answer = create_similarity_matrix_cosine(_utility_matrix)

    np.testing.assert_allclose(_student_answer, _solution, rtol=1e-1)
    print('success!')

    print(
        "Testing create_similarity_matrix_cosine on a dataset with a movie that has only been rated with 0: ",
        end='')

    _utility_matrix = DataFrame(
        [[5.0, nan, 4.5, nan, 5.0], [4.5, nan, 5.0, 1.5, 0.5],
         [0.0, 0.0, 0.0, 0.0, 0.0]], columns=[182, 195, 204, 376, 542],
        index=[2571, 2710, 4306])

    _solution = DataFrame([[1.0000000000, 0.8480436727, 0.000000000],
                           [0.8480436727, 1.0000000000, 0.0000000000],
                           [0.0000000000, 0.0000000000, 0.0000000000]],
                          columns=[2571, 2710, 4306], index=[2571, 2710, 4306])

    _student_answer = create_similarity_matrix_cosine(_utility_matrix)

    np.testing.assert_allclose(_student_answer, _solution, rtol=1e-1)
    print('success!')

def test_12(mean_center_columns, utility_matrix,
            create_similarity_matrix_cosine):
    print('Check output type: ', end='')
    _student_answer = mean_center_columns(utility_matrix)
    _expected_type = pd.DataFrame
    assert isinstance(_student_answer, _expected_type), f'expected output to be of type {_expected_type}, not {type(_solution)}'
    print('success!')

    print('Testing mean_center_columns on complete data set: ', end='')
    _solution1 = DataFrame([[0.5, 0.6666666666666665, 0.0, 0.5, 1.5],
                            [0.0, -1.3333333333333335, 0.5, -1.5, -3.0],
                            [-0.5, 0.6666666666666665, -0.5, 1.0, 1.5]],
                           columns=[182, 195, 204, 376, 542],
                           index=[2571, 2710, 4306])
    _student_answer = mean_center_columns(utility_matrix)

    np.testing.assert_allclose(_student_answer, _solution1, rtol=1e-1)
    print('success!')

    print('Testing create_similarity_matrix_cosine on complete data set: ',
          end='')
    _solution2 = DataFrame([[1.0, -0.9426042752501429, 0.8043933497362553],
                            [-0.9426042752501429, 1.0, -0.956600715515994],
                            [0.8043933497362553, -0.956600715515994, 1.0]],
                           columns=[2571, 2710, 4306],
                           index=[2571, 2710, 4306])
    _student_answer = create_similarity_matrix_cosine(_solution1)

    np.testing.assert_allclose(_student_answer, _solution2, rtol=1e-1)
    print('success!')

    print('Testing mean_center_columns on incomplete data set: ', end='')

    _utility_matrix = DataFrame(
        [[5.0, nan, 4.5, nan, 5.0], [4.5, nan, 5.0, 1.5, 0.5],
         [nan, 3.0, nan, nan, nan]], columns=[182, 195, 204, 376, 542],
        index=[2571, 2710, 4306])

    _solution1 = DataFrame(
        [[0.25, nan, -0.25, nan, 2.25], [-0.25, nan, 0.25, 0.0, -2.25],
         [nan, 0.0, nan, nan, nan]], columns=[182, 195, 204, 376, 542],
        index=[2571, 2710, 4306])
    _student_answer = mean_center_columns(_utility_matrix)

    np.testing.assert_allclose(_student_answer, _solution1, rtol=1e-1)
    print('success!')

    print('Testing create_similarity_matrix_cosine on incomplete data set: ',
          end='')

    _student_answer = create_similarity_matrix_cosine(_solution1)

    _solution2 = DataFrame([[1.0000000000, -0.9516522202, 0.0000000000],
                            [-0.9516522202, 1.0000000000, 0.0000000000],
                            [0.0000000000,  0.0000000000, 0.0000000000]],
        columns=[2571, 2710, 4306], index=[2571, 2710, 4306])

    np.testing.assert_allclose(_student_answer, _solution2, rtol=1e-1)
    print('success!')

def test_14(select_neighborhood, similarity, utility_matrix2):
    print('Check output type: ', end='')
    _solution = select_neighborhood(similarity[4306], utility_matrix2[123], 10)
    _expected_type = pd.Series
    assert isinstance(_solution,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(_solution)}'
    print('success!')

    print('Check select_neighborhood for user 123: ', end='')
    _solution = Series([0.4764040460, 0.2687109722, 0.569404],
                       index=[2571, 4444, 5555]).astype("float64")
    _student_answer = select_neighborhood(similarity[4306],
                                          utility_matrix2[123], 10)

    np.testing.assert_allclose(np.sort(_student_answer), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

    print('Check select_neighborhood for user 456: ', end='')
    _solution = Series([0.4764040460,  0.2687109722, 0.5694040620],
                       index=[2571, 4444, 5555]).astype("float64")
    _student_answer = select_neighborhood(similarity[4306],
                                          utility_matrix2[456], 10)

    np.testing.assert_allclose(np.sort(_student_answer), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

    print('Check select_neighborhood for movie 2710 and user 195: ', end='')
    _solution = Series([1.0], index=[2710]).astype("float64")
    _student_answer = select_neighborhood(similarity[2710],
                                          utility_matrix2[195], 10)

    np.testing.assert_allclose(np.sort(_student_answer), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

    print('Check empty neighborhood: ', end='')
    _dummy1 = pd.Series([np.nan, np.nan, 1.0, 1.0], index=[1, 2, 3, 4])
    _dummy2 = pd.Series([1.0, 1.0, np.nan, np.nan], index=[1, 2, 3, 4])
    _student_answer = select_neighborhood(_dummy1, _dummy2, 10)
    _solution = pd.Series([], index=[])
    np.testing.assert_allclose(np.sort(_student_answer), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

    print('Check only zero similarity: ', end='')
    _dummy1 = pd.Series([0.0, 0.0, 1.0, 1.0], index=[1, 2, 3, 4])
    _dummy2 = pd.Series([1.0, 1.0, np.nan, np.nan], index=[1, 2, 3, 4])
    _student_answer = select_neighborhood(_dummy1, _dummy2, 10)
    _solution = pd.Series([], index=[])
    np.testing.assert_allclose(np.sort(_student_answer), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

def test_15(select_neighborhood, similarity, utility_matrix2, weighted_mean):
    print('Check output type: ', end='')
    _neighborhood1 = select_neighborhood(similarity[4306],
                                         utility_matrix2[123], 10)
    _solution = weighted_mean(_neighborhood1, utility_matrix2[123])
    _expected_type = float
    assert isinstance(_solution,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(_solution)}'
    print('success!')

    print('Testing predictions: ', end='')

    # test prediction 1
    _neighborhood1 = select_neighborhood(similarity[4306],
                                         utility_matrix2[123], 10)
    _prediction1 = weighted_mean(_neighborhood1, utility_matrix2[123])
    assert _prediction1 > 4.1, "expected a predicted rating between 4.1 and 4.2"
    assert _prediction1 < 4.2, "expected a predicted rating between 4.1 and 4.2"

    # test prediction 2
    _neighborhood2 = select_neighborhood(similarity[4306],
                                         utility_matrix2[456], 10)
    _prediction2 = weighted_mean(_neighborhood2, utility_matrix2[456])
    assert _prediction2 > 2.2, "expected a predicted rating between 2.2 and 2.3"
    assert _prediction2 < 2.3, "expected a predicted rating between 2.2 and 2.3"
    print('success!')

    print('Testing prediction for empty neighborhood: ', end='')
    _empty = Series(0, index=[])
    _weighted_mean = weighted_mean(_empty, utility_matrix2[123])
    assert _weighted_mean is np.nan, "Expected: NaN"
    print('success!')

def test_17(utility_matrix3):
    print('Check output type: ', end='')
    _expected_type = pd.DataFrame
    assert isinstance(utility_matrix3,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(utility_matrix3)}'
    print('success!')

    print('Check select_neighborhood for user 123: ', end='')
    _solution = DataFrame(
        [[5.0, 1.5, nan, 4.0, 3.5], [5.0, 4.5, 4.0, 3.0, nan],
         [3.0, 1.0, 3.0, 2.0, nan], [4.5, 5.0, 4.0, 4.0, nan],
         [3.5, 1.5, 4.0, 4.5, nan], [1.5, 5.0, nan, 1.0, 3.5],
         [5.0, 0.5, 5.0, 4.5, nan]], columns=[2571, 2710, 4306, 4444, 5555],
        index=[123, 182, 195, 204, 376, 456, 542])

    np.testing.assert_allclose(np.sort(utility_matrix3), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

def test_19(centered_utility_matrix):
    print('Check output type: ', end='')
    _expected_type = pd.DataFrame
    assert isinstance(centered_utility_matrix,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(centered_utility_matrix)}'
    print('success!')

    print('Testing mean centered matrix on complete data set: ', end='')
    _solution1 = DataFrame(
        [[1.5, -2.0, nan, 0.5, 0.0], [0.875, 0.375, -0.125, -1.125, nan],
         [0.75, -1.25, 0.75, -0.25, nan], [0.125, 0.625, -0.375, -0.375, nan],
         [0.125, -1.875, 0.625, 1.125, nan], [-1.25, 2.25, nan, -1.75, 0.75],
         [1.25, -3.25, 1.25, 0.75, nan]],
        columns=[2571, 2710, 4306, 4444, 5555],
        index=[123, 182, 195, 204, 376, 456, 542])

    np.testing.assert_allclose(centered_utility_matrix, _solution1, rtol=1e-1)
    print('success!')

def test_20(similarity):
    print('Check output type: ', end='')
    _expected_type = pd.DataFrame
    assert isinstance(similarity,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(_solution)}'
    print('success!')

    print('Testing similarity of complete data set: ', end='')
    _solution1 = pd.DataFrame(
[[-0.8882176431559489, -0.5913123959890826, -2.7755575615628914e-17, 0.7749555280373117, 0.8278373543847156, 0.9091675835657295, 1.0],
[-0.575159934353966, -0.20149814827784024, -2.7755575615628914e-17, 0.1528941574312877, 0.3629750281770626, 0.6625413488689132, 1.0],
[-0.6363636363636362, -0.6239181323861147, 0.1528941574312877, 0.6949985619375698, 0.8278373543847156, 0.9185315234055025, 1.0],
[-0.9597599188661678, -0.8386592170224152, -0.6363636363636362, -0.5913123959890826, 0.6625413488689132, 0.7180944542557169, 1.0],
[-0.9597599188661678, -0.8699715932280792, -0.575159934353966, 0.6949985619375698, 0.7749555280373117, 0.9159337782355624, 1.0],
[-0.8882176431559489, -0.8699715932280792, -0.8429430938550735, -0.6239181323861147, 0.3629750281770626, 0.7180944542557169, 1.0],
[-0.8429430938550735, -0.8386592170224152, -0.20149814827784024, 0.9091675835657295, 0.9159337782355624, 0.9185315234055025, 1.0]],
            columns=[123, 182, 195, 204, 376, 456, 542],
            index=[123, 182, 195, 204, 376, 456, 542])

    np.testing.assert_allclose(np.sort(similarity), np.sort(_solution1), rtol=0.25)
    print('success!')

def test_21(neighborhood1, neighborhood2):
    print('Check output type: ', end='')
    _expected_type = pd.Series
    assert isinstance(neighborhood1,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(neighborhood1)}'
    print('success!')

    print('Check select_neighborhood for user 123: ', end='')
    _solution = Series(
        [0.774956, 0.827837, 0.909168],
        index=[195, 376, 542]).astype("float64")
    np.testing.assert_allclose(np.sort(neighborhood1), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

    print('Check select_neighborhood for user 456: ', end='')
    _solution = Series([0.3629750282, 0.7180944543],
                       index=[182, 204]).astype("float64")
    np.testing.assert_allclose(np.sort(neighborhood2), np.sort(_solution),
                               rtol=1e-1)
    print('success!')

def test_22(prediction1, prediction2):
    print('Check output type: ', end='')
    _expected_type = float
    assert isinstance(prediction1,
                      _expected_type), f'expected output to be of type {_expected_type}, not {type(_solution)}'
    print('success!')

    print("Testing predictions: ", end='')
    assert prediction1 > 3.9, "expected a predicted rating for user 123 above 3.9"
    assert prediction1 < 4.1, "expected a predicted rating for user 123 under 4.1"
    assert prediction2 > 3.9, "expected a predicted rating for user 456 above 3.9"
    assert prediction2 < 4.1, "expected a predicted rating for user 456 under 4.1"
    print('success!')

def test_24(utility, similarity):
    print('Check utility: ', end='')
    _partial_utility_solution = DataFrame(
        [[2.5, nan, nan, 4.0, nan, 2.5, 4.0, 2.5, nan, nan],
         [nan, nan, 3.5, nan, 4.0, 4.0, nan, nan, 2.0, nan],
         [3.5, 4.5, nan, 5.0, nan, nan, 4.0, nan, 4.5, 4.5],
         [nan, 3.5, nan, 4.0, nan, nan, 3.0, 4.0, nan, 4.5],
         [nan, 5.0, 4.5, nan, 4.5, 4.0, nan, 3.0, 4.5, 5.0]],
        columns=[68, 105, 177, 182, 232, 600, 603, 606, 608, 610],
        index=[1, 2, 16, 32, 47])
    _partial_student_solution = utility.loc[
        [1, 2, 16, 32, 47], [68, 105, 177, 182, 232, 600, 603, 606, 608, 610]]

    np.testing.assert_allclose(_partial_student_solution,
                               _partial_utility_solution, rtol=6e-1)
    print('success!')

    print('Check similarity: ', end='')
    _partial_similarity_solution = DataFrame(
        [[ 1.0000000000,  0.0057399843,  0.0173653612, -0.0260048241,  0.1013538293],
         [ 0.0057399843,  1.0000000000, -0.2764836030, -0.2554602920, -0.4963292866],
         [ 0.0173653612, -0.2764836030,  1.0000000000, -0.0558476293,  0.2085305218],
         [-0.0260048241, -0.2554602920, -0.0558476293,  1.0000000000,  0.4225094954],
         [ 0.1013538293, -0.4963292866,  0.2085305218,  0.4225094954,  1.0000000000]],
                                     columns=[1, 2, 16, 32, 47],
                                     index=[1, 2, 16, 32, 47])

    _partial_student_solution = similarity.loc[
        [1, 2, 16, 32, 47], [1, 2, 16, 32, 47]]

    np.testing.assert_allclose(_partial_student_solution,
                               _partial_similarity_solution, rtol=6e-1)
    print('success!')

def test_25(predict_ratings_item_based, similarity_items, utility_items, test_data):
    print('Computing solution: ', end='')
    _partial_solution = DataFrame(
        [[68, 2, 2.5, 3.4879827138689024], [68, 47, 4.0, 3.457850808877052],
         [68, 110, 2.5, 3.625071363263896], [68, 344, 2.5, 3.4579973887255604],
         [68, 593, 3.5, 3.4942987600576014]],
        columns=['userId', 'movieId', 'rating', 'predicted rating'],
        index=[1, 3, 6, 14, 27])
    _student_solution = predict_ratings_item_based(similarity_items,
                                                   utility_items, test_data[
                                                       ['userId', 'movieId',
                                                        'rating']]).head()

    print('success!')

    print('Testing layout of prediction: ', end='')
    for column in ['userId', 'movieId', 'rating', 'predicted rating']:
        assert column in _student_solution.columns, f'expected column {column} in output'

    print('success!')

    print('Testing values of prediction: ', end='')
    np.testing.assert_allclose(_student_solution, _partial_solution, rtol=2e-1)
    print('success!')

def test_26(mse, predicted_item_based):
    print('Testing mse item based: ', end='')
    np.testing.assert_allclose(mse(predicted_item_based), 0.6, rtol=0.15)
    print('success!')

def test_27(mse, predicted_user_based, predicted_item_based):
    print('Testing user based versus item based prediction: ', end='')
    assert mse(predicted_user_based) > mse(
        predicted_item_based), "expected a slightly lower error for item based prediction"
    print('success!')

def test_28(mse_random):
    print('Testing: ', end='')
    _student_solution = mse_random
    _solution = 3.2
    assert np.allclose(_student_solution, _solution,
                       atol=0.5), f'expected value around {_solution}'
    print('success!')

def test_29(mse_item_mean):
    print('Testing: ', end='')
    _student_solution = mse_item_mean
    _solution = 0.68
    assert np.allclose(_student_solution, _solution,
                       atol=0.2), f'expected value around {_solution}'
    print('success!')

def test_31(recommended, hidden, predicted_item_based, treshold_recommended):
    print('Testing: ', end='')
    tol = 5

    recommended_items_solution = 448
    recommended_items = recommended(predicted_item_based,
                                    treshold_recommended).shape[0]

    assert abs(recommended_items - recommended_items_solution) < tol, f'expected value around {recommended_items_solution}'

    hidden_items_solution = 545
    hidden_items = hidden(predicted_item_based, treshold_recommended).shape[0]

    assert abs(hidden_items - hidden_items_solution) < tol, f'expected value around {hidden_items_solution}'

    print('success!')

def test_32(used, unused, predicted_item_based, treshold_used):
    print('Testing: ', end='')
    tol = 5

    used_items_solution = 503
    used_items = used(predicted_item_based, treshold_used).shape[0]

    assert abs(used_items_solution - used_items) < tol, f'expected value around {used_items_solution}'

    unused_items_solution = 490
    unused_items = unused(predicted_item_based, treshold_used).shape[0]

    assert abs(unused_items_solution - unused_items) < tol, f'expected value around {unused_items_solution}'

    print('success!')

def test_33(confusion_matrix):
    print('Testing: ', end='')
    solution = DataFrame([[314, 134], [189, 356]], columns=['used', 'unused'],
                         index=['recommended', 'hidden'])

    pd.testing.assert_frame_equal(confusion_matrix, solution)
    print('success!')

def test_34(precision_item_based):
    print('Testing: ', end='')
    _student_solution = precision_item_based
    _solution = 0.7
    assert np.allclose(_student_solution, _solution,
                       atol=0.2), f'expected value around {_solution}'
    print('success!')

def test_35(recall_item_based):
    print('Testing: ', end='')
    _student_solution = recall_item_based
    _solution = 0.62
    assert np.allclose(_student_solution, _solution,
                       atol=0.2), f'expected value around {_solution}'
    print('success!')


# Errors:
# 1. if map_fn elements size is zero, then indexing using it's elements
# fails
# 2. Variables shape doesn't show change even after reassigning
# and make shape at creation time as validate_shape=False
# Change marg_vec_ to marg_vec_.read_value() in add marg, rem marg, add err, rem err ...
# 3. While loop frame error, if variables inside and outside while loop are used in
# same operation
# 4. Asychronous pains
# TF logical_and runs second terms even if first term fails
# Asynchronous prints
# 5. In while loop variables are not read after first run, i.e. other
# than loop vars no other ones run - https://github.com/tensorflow/tensorflow/issues/13616
# 6. tf cond expects same nested structure?? - How come this problem didn;t come up with
# other tf cond
# 7. https://github.com/tensorflow/tensor2tensor/issues/159

import tensorflow as tf

from collections import namedtuple

from data_loader import extract_data

C_SVM = 5.
RESERVE_THRESHOLD = -1

INF = 10**12 + 0.

import os
path_ = os.path.join("log_graph")

# Simple dot product - returns a rank 0 tensor/scalar
def simple_kernel(x1, x2):
  return tf.tensordot(x1, x2, [0,0])
  # val2 = tf.Print(val, [tf.constant("-----------Line 20"), val,
  #   tf.shape(val)])
  # return val2

model_params = {
  "C" : C_SVM,
  "eps" : RESERVE_THRESHOLD,
  "kernel" : simple_kernel,
  }

# Pipeline with fixed batch size of 1
def custom_input_fn(
  features, # Array like
  labels,
  ):
  assert features.shape[0] == labels.shape[0]
  feature_dict = {"x" : features}
  dataset = tf.data.Dataset.from_tensor_slices(
    (feature_dict, labels))
  return dataset.make_one_shot_iterator().get_next()

def calc_f(params, x_):
  x_all_ = tf.get_variable("x_all").read_value()
  y_all_ = tf.get_variable("y_all").read_value()
  b_ = tf.get_variable("b").read_value()
  marg_vec_ = tf.get_variable("marg_vec", dtype=tf.int32).read_value()
  err_vec_ = tf.get_variable("err_vec", dtype=tf.int32).read_value()
  alpha_marg_ = tf.get_variable("alpha_marg").read_value()

  alpha_err_ = tf.map_fn(lambda x: params["C"], err_vec_, dtype=tf.float32)
  # f(x_) = SUM(alpha[j] * y[j] * K(x_all[j] * x_)) + b

  # CHECK: Should I concat both the vectors to get all supp vectors
  # and then do the mulitplication (more parallelism) vs do them individually
  # and then join them (more distributed)??

  # More parallelism
  supp_vec_ = tf.concat([marg_vec_, err_vec_], 0)
  alpha_supp_ = tf.concat([alpha_marg_, alpha_err_], 0)

  # supp_vec_ = tf.Print(supp_vec_1, ["--------SUP VEC FOR f: ",supp_vec_1])
  temp = tf.cond(tf.equal(0, tf.shape(supp_vec_)[0]),
    lambda: tf.constant(0.),
    lambda: tf.map_fn(
      lambda i: y_all_[i] * params["kernel"](x_all_[i], x_),
      supp_vec_, dtype=tf.float32, name="58"))

  ans = tf.cond(tf.equal(0, tf.shape(supp_vec_)[0]),
    lambda: temp + b_,
    lambda: tf.tensordot(alpha_supp_, temp, [0,0], name="62") + b_)

  # ans = tf.Print(ans, [ans, temp, b_])
  return ans

# x_ - tensor
# y_ - tensor of rank 0/scalar
def calc_g(params, x_, y_):
  # g(x, y) = f(x)*y - 1
  return (calc_f(params, x_) * y_) - 1.

def rem_from_err(loc_index, loop_var):
  err_vec_ = loop_var.err_vec
  g_err_ = loop_var.g_err

  err_vec_ = tf.concat([err_vec_[:loc_index], err_vec_[loc_index+1:]], 0)
  # REmove g val of the vector
  g_err_ = tf.concat([g_err_[:loc_index], g_err_[loc_index+1:]], 0)
  # Force the above two operations
  return loop_var._replace(err_vec=err_vec_, g_err=g_err_)

def rem_from_rem(loc_index, loop_var):
  rem_vec_ = loop_var.rem_vec
  g_rem_ = loop_var.g_rem

  rem_vec_ = tf.concat([rem_vec_[:loc_index], rem_vec_[loc_index+1:]], 0)
  # REmove g val of the vector
  g_rem_ = tf.concat([g_rem_[:loc_index], g_rem_[loc_index+1:]], 0)
  # Force the above two operations
  return loop_var._replace(rem_vec=rem_vec_, g_rem=g_rem_)

def add_to_err(g_c, glob_index, loop_var):
  err_vec_ = loop_var.err_vec
  g_err_ = loop_var.g_err
  # Add new index to rem_vec
  err_vec_ = tf.concat([err_vec_, tf.reshape(glob_index, [1])], 0)

  # Add g_c to g_rem
  g_err_ = tf.concat([g_err_, tf.reshape(g_c, [1])], 0)

  loop_var = loop_var._replace(err_vec=err_vec_, g_err=g_err_)
  return loop_var

# Returns updated loop_var
# TODO: Add the reserve threshold check
def add_to_rem(eps, g_c, glob_index, loop_var=None):
  if loop_var is not None:
    rem_vec_ = loop_var.rem_vec
    g_rem_ = loop_var.g_rem
    # Add new index to rem_vec
    rem_vec_ = tf.concat([rem_vec_, tf.reshape(glob_index, [1])], 0)
    # Add g_c to g_rem
    g_rem_ = tf.concat([g_rem_, tf.reshape(g_c, [1])], 0)
    loop_var = loop_var._replace(rem_vec=rem_vec_, g_rem=g_rem_)
    return loop_var

  else:
    rem_vec_ = tf.get_variable("rem_vec", dtype=tf.int32)
    g_rem_ = tf.get_variable("g_rem")
    op1 = tf.assign(rem_vec_, tf.concat(
      [rem_vec_.read_value(), tf.reshape(glob_index, [1])], 0))
    # Add g_c to g_rem
    op2 = tf.assign(g_rem_, tf.concat(
      [g_rem_.read_value(), tf.reshape(g_c, [1])], 0))
    with tf.control_dependencies([op1, op2]):
      return tf.constant(1)

# Terminate if candidate g_c equals 0 or candidate alpha_c equals C
def termn_condn(C_svm, loop_var):
  print_op = tf.Print(tf.constant(0), [tf.constant("TREM CONDN"),
   loop_var.alpha_c, loop_var.g_c])
  with tf.control_dependencies([print_op]):
    return tf.logical_or(tf.equal(loop_var.g_c, 0.),
      tf.equal(loop_var.alpha_c, C_svm))

# TODO: do matrix inverse when size is one/inf
# Calculates the co-efficient sensitivities
def get_beta(kernel_fn, x_all_, y_all_, glob_index, marg_vec_, R_):
  # marg_vec_ = tf.get_variable("marg_vec", dtype=tf.int32).read_value()
  # R_ = tf.get_variable("R").read_value()
  # with tf.control_dependencies(None):
    x_c = x_all_[glob_index]
    y_c = y_all_[glob_index]

    # Calculating Q_(s_i, c) = y_all[s_i]*y_c*
    temp_q = tf.cond(tf.equal(tf.shape(marg_vec_)[0], 0),
      lambda: tf.constant([]),
      lambda: tf.map_fn(
              lambda i: y_c,
              # * y_all_[i] * kernel_fn(x_c, x_all_[i]),
              marg_vec_, dtype=tf.float32, name="BETA_MAP"),
      name="BETA_COND")
    # Add y_c as the first element
    temp_q = tf.concat([tf.reshape(y_c, [1,]), temp_q], 0)
    # reshape temp_q from (n,) to (n,1)/ rank 1 to rank 2 for matmul
    temp_q = tf.reshape(temp_q, [tf.shape(temp_q)[0], 1])
    # temp_q_p = tf.Print(temp_q, [tf.constant("LINE 150, BETA_CALC: "),
    #   tf.shape(x_all_), temp_q, R_])
    # print_op = tf.map_fn(lambda i: tf.Print(i, [i]), R_)
    # with tf.control_dependencies([print_op]):
    beta = -1 * tf.matmul(R_, temp_q)
    return tf.reshape(beta, [tf.shape(beta)[0]])

# Computes SUM(Q_ij*B_j)
def gamma_helper_fn(kernel_fn, beta_, indx, x_all_, y_all_, n, marg_vec_):
  # with tf.variable_scope(scope, reuse=True):

    iter_tensor = tf.range(tf.shape(marg_vec_)[0])
    # iter_tensor = tf.Print(iter_tensor_1, [tf.constant("GAMMA HELPER FN"),
      # tf.shape(x_all_), y_all_, n, indx, iter_tensor_1, ])
    return tf.cond(
      tf.equal(tf.shape(iter_tensor)[0], 0),
      lambda: tf.constant(0.),
      lambda: tf.reduce_sum(
        tf.map_fn(
          lambda i: ( y_all_[indx] * y_all_[marg_vec_[i]] * beta_[i + 1] *
            kernel_fn(x_all_[indx], x_all_[marg_vec_[i]]) ),
          iter_tensor,
          dtype=tf.float32
          )
        )
      )

# Making this more distributed
# vec is a tensor of indices for which gamma needs to be computed
def get_gamma(kernel_fn, vec, beta_, x_all_, y_all_, n, marg_vec_):
  x_c = x_all_[n-1]
  y_c = y_all_[n-1]
  # vec = tf.Print(vec_1, [tf.constant("---------------Line 174"), beta_,
  #   y_all_, tf.shape(y_all_), tf.shape(x_all_), tf.shape(beta_)])
  return tf.cond(tf.equal(tf.shape(vec)[0], 0),
    lambda: tf.constant([]),
    lambda: tf.map_fn(
      lambda i: (y_all_[i]*y_c*kernel_fn(x_all_[i], x_c) +
        beta_[0]*y_all_[i] +
        gamma_helper_fn(kernel_fn, beta_, i, x_all_, y_all_, n, marg_vec_)),
      vec,
      dtype=tf.float32)
    )

# TODO: Take beta as arg to prevent recalculation in case of new candidate vector
def add_to_marg_supp(kernel_fn, glob_index, gamma, alpha, x_all_,
  y_all_, loop_var):
  # with tf.variable_scope(scope, reuse=True):
  R_ = loop_var.R
  marg_vec_ = loop_var.marg_vec
  alpha_marg_ = loop_var.alpha_marg

  beta_indx = get_beta(kernel_fn, x_all_, y_all_, glob_index, marg_vec_, R_)
  beta_indx = tf.concat([beta_indx, [1]], 0)
  beta_shp = tf.shape(beta_indx)[0]

  beta_mat = tf.matmul(
    tf.reshape(beta_indx, [beta_shp, 1]),
    tf.reshape(beta_indx, [1, beta_shp]))

  beta_mat = 1/gamma * beta_mat

  # reshape R - will this be a problem while distributin???
  R_tens = R_
  R_tens = tf.concat([R_tens, tf.zeros([1, beta_shp-1])], 0)
  R_tens = tf.concat([R_tens, tf.zeros([beta_shp, 1])], 1)

  R_ = R_tens + beta_mat
  # add index to marg_vec_
  marg_vec_ = tf.concat([marg_vec_, tf.reshape(glob_index, [1])], 0)
  # add alpha to alpha_marg_
  alpha_marg_ = tf.concat([alpha_marg_, tf.reshape(alpha, [1])], 0),
  # Force the above two operations
  return loop_var._replace(R=R_, alpha_marg=alpha_marg_, marg_vec=marg_vec_)

# I need to test this function
def rem_helper(elem, R_, sz, k):
  i = elem//sz
  j = elem%sz
  return R_[i][j] - (1/R_[k][k])*R_[i][k]*R_[k][j]

# Need to assert that this value is same inverting Q_
def rem_from_marg_sup(kernel_fn, loc_index, loop_var):
  # with tf.variable_scope(scope, reuse=True):
  R_ = loop_var.R
  marg_vec_ = loop_var.marg_vec
  alpha_marg_ = loop_var.alpha_marg

  k = loc_index + 1

  sz = tf.shape(R_)[0]
  iter_tensor = tf.range(sz*sz)
  # iter_tensor = tf.reshape(iter_tensor, [sz, sz])

  # R_new = tf.Print(R_tens, [tf.constant("----------Line 252---"),
  #  tf.shape(R_tens), R_tens[0][0], sz])
  R_new = tf.map_fn(lambda i: rem_helper(i, R_, sz, k),
    iter_tensor, dtype=tf.float32)
  R_new = tf.reshape(R_new, [sz, sz])
  # # drop k th row
  R_new = tf.concat([R_new[:k, :], R_new[k+1:, :]], 0)
  # # drop kth col
  R_new = tf.concat([R_new[:, :k], R_new[:, k+1:]], 1)

  R_ = R_new

  # remove index from marg_vec_
  marg_vec_ = tf.concat([marg_vec_[:loc_index], marg_vec_[loc_index+1:]], 0)
  # remove alpha value from alpha_marg
  alpha_marg_= tf.concat(
    [alpha_marg_[:loc_index], alpha_marg_[loc_index+1:]], 0)
  # Force the update operations
  return loop_var._replace(R=R_, marg_vec=marg_vec_, alpha_marg=alpha_marg_)

# Returns a loop_var of the updated loop_var
def handle_marg(params, min_indx, beta_, loop_var):
  # with tf.variable_scope(scope, reuse=True):
  marg_vec_ = loop_var.marg_vec
  # If beta[min_indx+1] < 0, then alpha becomes 0 and joins remaining vec,
  # else alpha becomes C and joins error vectors
  loop_var = tf.cond(
        beta_[min_indx+1] < 0,
        lambda: add_to_rem(params["eps"],
          tf.constant(0.), marg_vec_[min_indx], loop_var),
        lambda: add_to_err(tf.constant(0.), marg_vec_[min_indx], loop_var)
        )
  # Remove vector from R_ and marg_vec_
  # with tf.control_dependencies([op1]):
  return rem_from_marg_sup(params["kernel"], min_indx, loop_var)

def update_g(g_tens, gamma_tens, min_alpha):
  iter_tensor = tf.range(tf.shape(g_tens)[0])
  return tf.cond(
    tf.equal(tf.shape(g_tens)[0], 0),
    lambda: tf.constant([]),
    lambda: tf.map_fn(
      lambda i: g_tens[i] + min_alpha*gamma_tens[i],
      iter_tensor, dtype=tf.float32)
  )

# Tensor should have non-zero shape
def check_min(tens, val):
  min_indx = tf.argmin(tens)
  return tf.cond(tf.equal(tens[min_indx], val),
    lambda: tf.cast(min_indx, tf.int32),
    lambda: tf.shape(tens)[0] + 1)

# returns alpha_c and g_c
def mini_iter(params, iter_ct, loop_var):
  # with tf.control_dependencies(None):
    marg_vec_ = loop_var.marg_vec
    err_vec_ = loop_var.err_vec
    rem_vec_ = loop_var.rem_vec
    alpha_marg_ = loop_var.alpha_marg
    g_err_ = loop_var.g_err
    g_rem_ = loop_var.g_rem
    g_c_ = loop_var.g_c
    alpha_c_ = loop_var.alpha_c
    b_ = loop_var.b
    R_ = loop_var.R

    x_all_ = loop_var.x_all
    y_all_ = loop_var.y_all
    n = loop_var.n

    printer = tf.Print(tf.constant(0), [tf.constant("AAAAAAAAAAAA")])
    # Handle case of matrix containing inf
    R_ = tf.cond(
      tf.equal(tf.shape(loop_var.marg_vec)[0], 0),
      lambda: tf.constant([[-INF]]),
      lambda: loop_var.R)
    loop_var = loop_var._replace(R=R_)

    print_op = tf.Print(tf.constant(1.), [tf.constant("----------Line 348 ------: "),
      tf.constant("MARG: "), loop_var.marg_vec,
      tf.constant("ERR: "), loop_var.err_vec,
      tf.constant("REM: "), loop_var.rem_vec,
      tf.constant("ALPHA: "), loop_var.alpha_marg,
      tf.constant("B: "), loop_var.b,
      tf.constant("R: "), loop_var.R,
      iter_ct], name="BITCH_PRINT")

    # Get beta and gamma necessary for the iteration
    # with tf.control_dependencies([print_op]):
    beta_ = get_beta(params["kernel"], x_all_, y_all_,
      n-1, marg_vec_, R_)
      # Gamma for err_vec
    gamma_err_ = get_gamma(params["kernel"], err_vec_, beta_,
      x_all_, y_all_, n, marg_vec_)
    # Gamma for rem_vec
    gamma_rem_ = get_gamma(params["kernel"], rem_vec_, beta_,
      x_all_, y_all_, n, marg_vec_)
    # Gamma for candidate vector
    n_temp = n-1
    gamma_c_ = get_gamma(params["kernel"], tf.reshape(n_temp, [1]), beta_,
      x_all_, y_all_, n, marg_vec_)

    with tf.control_dependencies([beta_]):
      return (iter_ct+1, loop_var)

    # # Book-keeping -
    # # For margin support vectors, if Beta_s > 0, alpha_c can go to C or
    # # if Beta_s < 0, alpha_s can go to 0, causing a transition in state
    # iter_tensor_a = tf.range(tf.shape(marg_vec_)[0])
    # # Add extra condition to prevent map_fn on empty iter_tensor
    # marg_trans_alpha = tf.cond(tf.equal(tf.shape(iter_tensor_a)[0], 0),
    #   lambda: tf.constant([]),
    #   lambda: tf.map_fn(
    #     lambda i: tf.cond(beta_[i+1] < 0,
    #       # Alpha goes to zero
    #       lambda: -1 * alpha_marg_[i] / beta_[i+1],
    #       lambda: tf.cond(beta_[i+1] > 0,
    #         # Alpha goes to C
    #         lambda: (params["C"] - alpha_marg_[i]) / beta_[i+1],
    #         # Beta is zero
    #         lambda: tf.constant(float("Inf"))
    #         )
    #       ),
    #     iter_tensor_a, dtype=tf.float32)
    # )
    # # For error support vectors, if Gamma_i > 0, then g_i increases to 0
    # # and causes state change
    # iter_tensor_e = tf.range(tf.shape(err_vec_)[0])
    # err_trans_alpha = tf.cond(tf.equal(tf.shape(iter_tensor_e)[0], 0),
    #   lambda: tf.constant([]),
    #   lambda: tf.map_fn(
    #     lambda i: tf.cond(gamma_err_[i] > 0,
    #       lambda: -1 * g_err_[i] / gamma_err_[i],
    #       lambda: tf.constant(float("Inf"))
    #       ),
    #     iter_tensor_e, dtype=tf.float32)
    #   )

    # # For the remaining vectors, if Gamma_i < 0, then g_i decreases to 0
    # iter_tensor_r = tf.range(tf.shape(rem_vec_)[0])
    # rem_trans_alpha = tf.cond(tf.equal(tf.shape(iter_tensor_r)[0], 0),
    #   lambda: tf.constant([]),
    #   lambda: tf.map_fn(
    #     lambda i: tf.cond(gamma_rem_[i] < 0,
    #       lambda: -1 * g_rem_[i] / gamma_rem_[i],
    #       lambda: tf.constant(float("Inf"))
    #       ),
    #     iter_tensor_r, dtype=tf.float32)
    #   )
    # # For candidate vector, if gamma_i > 0, then g_i increases to 0
    # # or alpha_c increases to C_svm
    # candidate_trans_alpha = tf.minimum(
    #   tf.cond(gamma_c_[0] > 0,
    #     lambda: -1*g_c_/gamma_c_[0],
    #     lambda: tf.constant(float("Inf"))),
    #   params["C"] - alpha_c_)

    # # Compare the min and do the necessary transitions
    # all_trans_alpha = tf.concat(
    #   [marg_trans_alpha, err_trans_alpha, rem_trans_alpha,
    #   tf.reshape(candidate_trans_alpha, [1,])], 0)
    # min_alpha = tf.reduce_min(all_trans_alpha)

    # # CALCULATING NEW VALUES
    # # Calculate new b
    # b_ = b_ + beta_[0]*min_alpha
    # loop_var = loop_var._replace(b=b_)
    # # Calculate new values of alpha for marg vec
    # iter_tensor = tf.range(tf.shape(alpha_marg_)[0])
    # # Condition to not run the map fn on zero sz tensor
    # alpha_marg_ = tf.cond(tf.equal(tf.shape(iter_tensor)[0], 0),
    #   lambda: tf.constant([]),
    #   lambda: tf.map_fn(
    #     lambda i: alpha_marg_[i] + beta_[i+1]*min_alpha,
    #     iter_tensor, dtype=tf.float32)
    #   )
    # loop_var = loop_var._replace(alpha_marg=alpha_marg_)
    # # Calculate new value of alpha_c
    # alpha_c_ = alpha_c_ + min_alpha
    # loop_var = loop_var._replace(alpha_c=alpha_c_)

    # # Calculate new values of g_all
    # g_err_ = update_g(g_err_, gamma_err_, min_alpha)
    # g_rem_ = update_g(g_rem_, gamma_rem_, min_alpha)
    # g_c_ = g_c_ + gamma_c_[0]*min_alpha
    # loop_var = loop_var._replace(g_rem=g_rem_, g_err=g_err_, g_c=g_c_)

    # print_op2 = tf.Print(tf.constant(1.),
    #   [tf.constant("BETA: "), beta_,
    #   tf.constant("GAMMA_c: "), gamma_c_,
    #   tf.constant("ALL TRANS: "), all_trans_alpha])
    # with tf.control_dependencies([print_op2]):
    #   # MOVING VECTORS
    #   # Check if min is in marg sup vectors
    #   # Run get index only if tnsor sz is non-zero
    #   min_indx = tf.cond(tf.equal(tf.shape(marg_trans_alpha)[0], 0),
    #     lambda: tf.constant(1),
    #     lambda: check_min(marg_trans_alpha, min_alpha),
    #     )
    #   loop_var = tf.cond(
    #     min_indx < tf.shape(marg_vec_)[0],
    #     lambda: handle_marg(params, min_indx, beta_, loop_var),
    #     lambda: loop_var,
    #     name="476")

    #   # # Check if min is in err sup vectors
    #   # def handle_err(loc_index, loop_var):
    #   #   loop_var = add_to_marg_supp(params["kernel"], err_vec_[loc_index],
    #   #     gamma_err_[loc_index], params["C"], x_all_, y_all_, loop_var)
    #   #   loop_var = rem_from_err(loc_index, loop_var)
    #   #   return loop_var

    #   # min_indx = tf.cond(tf.equal(tf.shape(err_trans_alpha)[0], 0),
    #   #   lambda: tf.constant(1),
    #   #   lambda: check_min(err_trans_alpha, min_alpha)
    #   #   )

    #   # loop_var = tf.cond(
    #   #   min_indx < tf.shape(err_trans_alpha)[0],
    #   #   lambda: handle_err(min_indx, loop_var),
    #   #   lambda: loop_var,)

    #   # # Check if min is in remaining vectors
    #   # def handle_rem(loc_index, loop_var):
    #   #   loop_var = add_to_marg_supp(params["kernel"], rem_vec_[loc_index],
    #   #     gamma_rem_[loc_index], 0., x_all_, y_all_, loop_var)
    #   #   loop_var = rem_from_rem(loc_index, loop_var)
    #   #   return loop_var

    #   # min_indx = tf.cond(tf.equal(tf.shape(rem_trans_alpha)[0], 0),
    #   #   lambda: tf.constant(1),
    #   #   lambda: check_min(rem_trans_alpha, min_alpha)
    #   #   )
    #   # loop_var = tf.cond(
    #   #   min_indx < tf.shape(rem_trans_alpha)[0],
    #   #   lambda: handle_rem(min_indx, loop_var),
    #   #   lambda: loop_var)
    #   # Need to add the candidate vector here as I need g and gamma values
    #   # loop_var = tf.cond(
    #   #   termn_condn(params["C"], loop_var),
    #   #   lambda: tf.cond(
    #   #     tf.equal(loop_var.g_c, 0.),
    #   #     lambda: add_to_marg_supp(params["kernel"],
    #   #       n-1, gamma_c_[0], loop_var.alpha_c, x_all_, y_all_, loop_var),
    #   #     lambda: add_to_err(loop_var.g_c, n-1, loop_var)),
    #   #   lambda: loop_var)

    #   printer = tf.Print(tf.constant(1.),
    #       [tf.constant("----------------Line 512----: "), n, loop_var.marg_vec,
    #       loop_var.err_vec, loop_var.rem_vec, loop_var.alpha_marg, loop_var.g_err,
    #       loop_var.R, loop_var.b])
    #   with tf.control_dependencies([printer]):
    #     return loop_var

def fit_point(params):
  marg_vec_ = tf.get_variable("marg_vec", dtype=tf.int32)
  err_vec_ = tf.get_variable("err_vec", dtype=tf.int32)
  rem_vec_ = tf.get_variable("rem_vec", dtype=tf.int32)
  alpha_marg_ = tf.get_variable("alpha_marg")
  g_err_ = tf.get_variable("g_err")
  g_rem_ = tf.get_variable("g_rem")
  g_c_ = tf.get_variable("g_c")
  alpha_c_ = tf.get_variable("alpha_c")
  b_ = tf.get_variable("b")
  R_ = tf.get_variable("R")

  # While loop supports named tuples and not dicts
  LoopVars = namedtuple("LoopVars",
    "x_all, y_all, n, marg_vec, err_vec, rem_vec, alpha_marg, g_err, g_rem, g_c, alpha_c, b, R")
  init_val = (tf.constant(0), LoopVars(
    x_all=tf.identity(tf.get_variable("x_all").read_value()) + 0.,
    y_all=tf.identity(tf.get_variable("y_all").read_value()) + 0.,
    n=tf.identity(tf.get_variable("n", dtype=tf.int32).read_value()) + 0,
    marg_vec=marg_vec_.read_value(),
    err_vec=err_vec_.read_value(),
    rem_vec=rem_vec_.read_value(),
    alpha_marg=alpha_marg_.read_value(),
    g_err=g_err_.read_value(),
    g_rem=g_rem_.read_value(),
    g_c=g_c_.read_value(),
    alpha_c=alpha_c_.read_value(),
    b=b_.read_value(),
    R=R_.read_value(),
    ))

  # init_val = tf.Print(init_val, [tf.constant("----------------Line 509--- "),
  #   fixed_vars["n"], fixed_vars["x_all"], fixed_vars["y_all"],
  #   tf.constant("------G_c: "), init_val])
  fin_i, fin_vars = tf.while_loop(
    cond=lambda i, p: tf.logical_and(tf.logical_not(termn_condn(params["C"], p)),
      tf.less(i, 1)),
    body=lambda i, p: mini_iter(params, i, p),
    loop_vars=init_val,
    parallel_iterations=1,
    back_prop=False,
    name="FIT_WHILE"
    )

  upd_marg_vec = tf.assign(marg_vec_, fin_vars.marg_vec, validate_shape=False)
  upd_err_vec = tf.assign(err_vec_, fin_vars.err_vec, validate_shape=False)
  upd_rem_vec = tf.assign(rem_vec_, fin_vars.rem_vec, validate_shape=False)
  upd_alpha_marg = tf.assign(alpha_marg_, fin_vars.alpha_marg,
    validate_shape=False)
  upd_g_err = tf.assign(g_err_, fin_vars.g_err, validate_shape=False)
  upd_g_rem = tf.assign(g_rem_, fin_vars.g_rem, validate_shape=False)
  upd_alpha_c = tf.assign(alpha_c_, fin_vars.alpha_c, validate_shape=False)
  upd_g_c = tf.assign(g_c_, fin_vars.g_c, validate_shape=False)
  upd_b = tf.assign(b_, fin_vars.b, validate_shape=False)
  upd_R = tf.assign(R_, fin_vars.R, validate_shape=False)

  with tf.control_dependencies([upd_marg_vec, upd_err_vec, upd_rem_vec,
    upd_alpha_marg, upd_g_err, upd_g_rem, upd_alpha_c, upd_g_c, upd_b, upd_R]):
    return fin_i

def create_svm_variables(x_shape):
  with tf.variable_scope("svm_model") as scope:
    # Variable creation
    # count of the number of data points seen
    n_ = tf.get_variable("n", initializer=tf.constant(0), trainable=False,
      dtype=tf.int32)
    # offset
    b_ = tf.get_variable("b", initializer=tf.constant(0.), trainable=False,)
    # Variable tensor - containing all the data points and labels
    x_all_ = tf.get_variable("x_all", shape=[0, x_shape],
      validate_shape=False)
    y_all_ = tf.get_variable("y_all", shape=[0],
      validate_shape=False)

    # Variable tensor - representing the Margin Support vector indices
    marg_vec_ = tf.get_variable("marg_vec", shape=[0], dtype=tf.int32,
      validate_shape=False)
    # Variable tensor - representing the Error Support vector indices
    err_vec_ = tf.get_variable("err_vec", shape=[0], dtype=tf.int32,
      validate_shape=False)
    # Variable tensor - representing the Remaining vector indices
    rem_vec_ = tf.get_variable("rem_vec", shape=[0], dtype=tf.int32,
      validate_shape=False)

    # Variable for alpha of margin vectors
    alpha_marg_ = tf.get_variable("alpha_marg", shape=[0],
      validate_shape=False)
    # Variable denoting g of error support vectors
    g_err_ = tf.get_variable("g_err", shape=[0],
      validate_shape=False)
    # Variable denoting g of remaining vectors
    g_rem_ = tf.get_variable("g_rem", shape=[0],
      validate_shape=False)

    # Variable for the inverse Jacobian matrix R - initially Inf
    R_ = tf.get_variable("R", initializer=tf.constant([[INF]]), dtype=tf.float32,
      validate_shape=False)

    # Variable denoting the alpha of new candidate
    alpha_c_ = tf.get_variable("alpha_c", initializer=0.)
    # Variable denoting the g of candidate vector
    g_c_ = tf.get_variable("g_c", shape=())

    return scope

def return_w_b(scope, params):
  # get a1 and a2 for a1x+a2y+b=0
  with tf.variable_scope(scope, reuse=True):
    marg_vec_ = tf.get_variable("marg_vec")
    err_vec_ = tf.get_variable("err_vec")
    alpha_marg_ = tf.get_variable("alpha_marg")
    x_all_ = tf.get_variable("x_all")
    y_all_ = tf.get_variable("y_all")
    b_ = tf.get_variable("b")
    n_ = tf.get_variable("n")

    w = tf.zeros([2])

    iter_tensor = tf.range(tf.shape(marg_vec_)[0])
    temp_marg = tf.map_fn(
      lambda i: y_all_[marg_vec_[i]] * x_all_[marg_vec_[i]] * alpha_marg_[i],
      iter_tensor)

    iter_tensor = tf.range(tf.shape(err_vec_)[0])
    temp_err = tf.map_fn(
      lambda i: y_all_[err_vec_[i]] * x_all_[err_vec_[i]] * params["C"],
      iter_tensor)

    w = tf.reduce_sum(tf.concat([temp_marg, temp_err], 0), 0)
    return (w, b_, x_all_, y_all_, n_)

def draw_dec_bdry(w, b, x_all, y_all, n):
  import numpy as np

  from matplotlib import colors
  from matplotlib import pyplot as plt

  x_all = np.array(x_all)
  y_all = np.array(y_all)

  plt.scatter(x_all[:n, 0], x_all[:n, 1], c=y_all[:n],
    cmap=colors.ListedColormap(['orange', 'red', 'green']))
  plt.plot()

  start = 0
  last = 1
  xl = [start, last]
  yl = []
  for xx in xl:
    yl.append((-b - w[0]*xx)/w[1])
  plt.plot(xl, yl)
  plt.show()

def svm_train(x_, y_, params):
  x_all_ = tf.get_variable("x_all")
  y_all_ = tf.get_variable("y_all")
  n_ = tf.get_variable("n", dtype=tf.int32)
  alpha_c_ = tf.get_variable("alpha_c")
  g_c_ = tf.get_variable("g_c")

  #  Concatanate x_all, y_all to include the new vector and it's labe
  op1 = tf.assign(x_all_,
    tf.concat([x_all_, tf.reshape(x_, [1, params["shape"]])], 0),
    validate_shape=False)
  op2 = tf.assign(y_all_,
    tf.concat([y_all_, tf.reshape(y_, [1])], 0),
    validate_shape=False, name="y_assign531")
  # Update count of data points seen
  op3 = tf.assign_add(n_, 1)
  # Reset alpha_c
  op4 = tf.assign(alpha_c_, 0.)
  # Force the above operations
  with tf.control_dependencies([op1, op2, op3, op4]):
    # g_c = tf.assign(g_c_, 0)
    g_c = tf.assign(g_c_, calc_g(params, x_, y_),
        validate_shape=False)
    return tf.cond(
      tf.greater(g_c, 0),
      lambda: add_to_rem(params["eps"], g_c, n_.read_value()-1),
      lambda: fit_point(params))

def svm_model_fn():
  train_file = "data_1.csv"
  x_train, y_train = extract_data(train_file, "csv")
  model_params["shape"] = x_train.shape[1]

  x_ = tf.placeholder(tf.float32, shape=(x_train.shape[1]))
  y_ = tf.placeholder(tf.float32, shape=())
  scope = create_svm_variables(x_shape=x_train.shape[1])

  with tf.variable_scope(scope, reuse=True):
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      tr = svm_train(x_, y_, model_params)
      for i in range(1):
        print("\n \n")
        print(sess.run(tr, feed_dict={x_ : x_train[i], y_ : y_train[i]}))
        # y_all_ = tf.get_variable("y_all")
        # print(sess.run(y_all_))
      # writer = tf.summary.FileWriter(path_, sess.graph)
      # writer.close()
        # print(sess.run(r[0]))

svm_model_fn()

# class Inc_SVM():
#   def __init__(train_file, C_svm, eps_thresh, kernel_fn):
#     self.params = {"C" : C_svm, "eps": eps_thresh, "kernel": kernel_fn}


#pip install -q -U "tensorflow-gpu==2.0.0b1"
#pip install -q -U tensorflow_hub
import warnings
warnings.filterwarnings('ignore')

import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
import tensorflow_hub as hub
import logging
import argparse
import sys
import json
from PIL import Image

parser = argparse.ArgumentParser ()
parser.add_argument ('image_dir')
parser.add_argument('checkpoint')
parser.add_argument ('--top_k', default = 5, type = int)
parser.add_argument ('--category_names', default = 'label_map.json')

commands = parser.parse_args()
image_path = commands.image_dir
export_path_keras = commands.checkpoint
classes = commands.category_names
top_k = commands.top_k
reloaded = tf.keras.models.load_model(export_path_keras, custom_objects={'KerasLayer': hub.KerasLayer})

with open(classes, 'r') as f:
    class_names = json.load(f)

def process_image(image):
    image = tf.convert_to_tensor(image, dtype=tf.float32)
    image = tf.image.resize(image, (224, 224))
    image /= 255
    return image.numpy()

def predict(image_path, model, top_k):
    im = Image.open(image_path)
    im = np.asarray(im)
    im = process_image(im)
    expand_im = np.expand_dims(im, axis=0)
    p = model.predict(expand_im)
    top_k_values, top_k_indices = tf.nn.top_k(p, k=top_k)

    top_k_values = top_k_values.numpy()
    top_k_indices = top_k_indices.numpy()
    flower_classes = []
    for idx in top_k_indices[0]:
        flower_classes.append(class_names[str(idx+1)])
    return top_k_values, top_k_indices, flower_classes

probs, classes, flowers = predict(image_path, reloaded, top_k)
print("Possibility: ",probs)
print("Class: ",classes)
print("Flower: ",flowers)


class UnionFind:
    def __init__(self):
        self._parent = {}
        self._size = {}

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a == b:
            return
        if self._size[a] < self._size[b]:
            a, b = b, a
        self._parent[b] = a
        self._size[a] += self._size[b]

    def find(self, x):
        if x not in self._parent:
            self._parent[x] = x
            self._size[x] = 1
        while self._parent[x] != x:
            self._parent[x] = self._parent[self._parent[x]]
            x = self._parent[x]
        return x

class Solution:
    def smallestEquivalentString(self, s1: str, s2: str, baseStr: str) -> str:
        uf = UnionFind()

        for a, b in zip(s1, s2):
            uf.union(a, b)

        parents = defaultdict(list)

        for x in s1 + s2:
            parent = uf.find(x)
            parents[parent].append(x)

        for key in parents.keys():
            parents[key].sort()

        res = []

        for x in baseStr:
            parent = uf.find(x)
            if parent not in parents:
                res.append(x)
            else:
                res.append(parents[parent][0])

        return "".join(res)


class node:
    def __init__(self,value):
        self.right=None
        self.left=None
        self.value=value

def balanced(b):
    maxh=0
    minh=-1
    l=[]
    current=b
    currenth=0
    while current!=None or len(l)>0:
        if current!=None:
            l.append(current)
            current=current.left
            currenth+=1
        else:
            current=l.pop()
            currenth-=1
            if current.right==None and current.left==None:
                maxh=max(maxh,currenth)
                if minh==-1: minh=currenth
                else: minh=min(minh,currenth)
            current=current.right
    return maxh-minh<2