require 'uri'
require 'cgi'

require 'ghtorrent/api_client'
require 'ghtorrent/settings'
require 'ghtorrent/utils'
require 'ghtorrent/logging'

module GHTorrent
  module Retriever

    include GHTorrent::Settings
    include GHTorrent::Utils
    include GHTorrent::APIClient
    include GHTorrent::Logging

    def persister
      raise Exception.new("Unimplemented")
    end

    def retrieve_user_byusername(user)
      stored_user = persister.find(:users, {'login' => user})
      if stored_user.empty?
        url = ghurl "users/#{user}"
        u = api_request(url)

        if u.nil? or u.empty?
          return
        end

        unq = persister.store(:users, u)
        what = user_type(u['type'])
        info "Added user #{what} #{user}"
        u
      else
        what = user_type(stored_user.first['type'])
        debug "#{what} #{user} exists"
        stored_user.first
      end
    end

    # Try Github user search by email. This is optional info, so
    # it may not return any data. If this fails, try searching by name
    # http://developer.github.com/v3/search/#email-search
    def retrieve_user_byemail(email, name)
      url = ghurl("legacy/user/email/#{CGI.escape(email)}")
      byemail = api_request(url)

      if byemail.nil? or byemail.empty?
        # Only search by name if name param looks like a proper name
        byname = if not name.nil? and name.split(/ /).size > 1
                  url = ghurl("legacy/user/search/#{CGI.escape(name)}")
                  api_request(url)
                 end

        if byname.nil? or byname['users'].nil? or byname['users'].empty?
          nil
        else
          user = byname['users'].find do |u|
                u['name'] == name and
                not u['login'].nil? and
                not retrieve_user_byusername(u['login']).nil?
          end

          unless user.nil?
            # Make extra sure that if we got an email it matches that
            # of the retrieved user
            if not email.nil? and user['email'] == email
              user
            else
              warn "Could not find user #{email}"
              nil
            end
          else
            warn "Could not find user #{email}"
            nil
          end
        end
      else
        unless byemail['user']['login'].nil?
          info "Added user #{byemail['user']['login']} retrieved by email #{email}"
          retrieve_user_byusername(byemail['user']['login'])
        else
          u = byemail['user']
          unq = persister.store(:users, u)
          what = user_type(u['type'])
          info "Added user #{what} #{user}"
          u
        end
      end
    end

    def retrieve_user_follower(followed, follower)
      stored_item = persister.find(:followers, {'follows' => followed,
                                                'login' => follower})

      if stored_item.empty?
        retrieve_user_followers(followed).find{|x| x['login'] == follower}
      else
        stored_item.first
      end
    end

    def retrieve_user_followers(user)
      followers = paged_api_request(ghurl "users/#{user}/followers")
      followers.each do |x|
        x['follows'] = user

        exists = !persister.find(:followers, {'follows' => user,
                                              'login' => x['login']}).empty?

        if not exists
          persister.store(:followers, x)
          info "Added follower #{user} -> #{x['login']}"
        else
          debug "Follower #{user} -> #{x['login']} exists"
        end
      end

      persister.find(:followers, {'follows' => user})
    end

    def retrieve_user_following(user)
      following = paged_api_request(ghurl "users/#{user}/following")
      user_followers_entry = nil

      following.each do |x|

        if user_followers_entry.nil?
          reverse_lookup = persister.find(:followers, {'follows' => x['login'],
                                                       'login' => user})
          if reverse_lookup.empty?
            user_followers_entry = retrieve_user_followers(x['login']).\
                                     find{|y| y['login'] == user}
          else
            user_followers_entry = reverse_lookup[0]
          end
        end

        exists = !persister.find(:followers, {'follows' => x['login'],
                                              'login' => user}).empty?
        if not exists
          user_followers_entry['follows'] = x['login']
          user_followers_entry.delete(:_id)
          user_followers_entry.delete('_id')
          a = persister.store(:followers, user_followers_entry)
          info "Added following #{user} -> #{x['login']}"
        else
          debug "Following #{user} -> #{x['login']} exists"
        end
      end

      persister.find(:followers, {'login' => user})
    end

    def retrieve_pull_request_commit(pr_obj, repo, sha, user)
      pull_commit = persister.find(:pull_request_commits, {'sha' => "#{sha}"})

      if pull_commit.empty?
        commit = retrieve_commit(repo, sha, user)

        unless commit.nil?
          commit.delete(nil)
          commit.delete(:_id)
          commit.delete('_id')
          commit['pull_request_id'] = pr_obj[:id]
          persister.store(:pull_request_commits, commit)
          info "Added commit #{user}/#{repo} -> #{sha} with pull_id #{pr_obj['id']}"
          commit
        else
          return
        end
      else
        debug "Pull request commit #{user}/#{repo} -> #{sha} exists for pull id #{pr_obj['id']}"
        pull_commit.first
      end
    end

    # Retrieve a single commit from a repo
    def retrieve_commit(repo, sha, user)
      commit = persister.find(:commits, {'sha' => "#{sha}"})

      if commit.empty?
        url = ghurl "repos/#{user}/#{repo}/commits/#{sha}"
        c = api_request(url)

        if c.nil? or c.empty?
          return
        end

        # commit patches are big and not always interesting
        if config(:commit_handling) == 'trim'
          c['files'].each { |file| file.delete('patch') }
        end
        persister.store(:commits, c)
        info "Added commit #{user}/#{repo} -> #{sha}"
        c
      else
        debug "Commit #{user}/#{repo} -> #{sha} exists"
        commit.first
      end
    end

    # Retrieve commits starting from the provided +sha+
    def retrieve_commits(repo, sha, user, pages = -1)

      url = if sha.nil?
              ghurl "repos/#{user}/#{repo}/commits"
            else
              ghurl "repos/#{user}/#{repo}/commits?sha=#{sha}"
            end

      commits = restricted_page_request(url, pages)

      commits.map do |c|
        retrieve_commit(repo, c['sha'], user)
      end.select{|x| not x.nil?}
    end

    def retrieve_repo(user, repo, refresh = false)
      stored_repo = persister.find(:repos, {'owner.login' => user,
                                             'name' => repo })
      if stored_repo.empty? or refresh
        url = ghurl "repos/#{user}/#{repo}"
        r = api_request(url)

        if r.nil? or r.empty?
          return
        end

        if refresh
          persister.upsert(:repos, {'name' => r['name'], 'owner.login' => r['owner']['login']}, r)
          info "Refreshed repo #{user} -> #{repo}"
        else
          persister.store(:repos, r)
          info "Added repo #{user} -> #{repo}"
        end
        r
      else
        debug "Repo #{user} -> #{repo} exists"
        stored_repo.first
      end
    end

    def retrieve_languages(owner, repo)
      paged_api_request ghurl "repos/#{owner}/#{repo}/languages"
    end

    # Retrieve organizations the provided user participates into
    def retrieve_orgs(user)
      url = ghurl "users/#{user}/orgs"
      orgs = paged_api_request(url)
      orgs.map{|o| retrieve_org(o['login'])}
    end

    # Retrieve a single organization
    def retrieve_org(org)
      retrieve_user_byusername(org)
    end

    # Retrieve organization members
    def retrieve_org_members(org)
      stored_org_members = persister.find(:org_members, {'org' => org})

      org_members = paged_api_request(ghurl "orgs/#{org}/members")
      org_members.each do |x|
        x['org'] = org

        exists = !stored_org_members.find { |f|
          f['org'] == org && f['login'] == x['login']
        }.nil?

        if not exists
          persister.store(:org_members, x)
          info "Added org_member #{org} -> #{x['login']}"
        else
          debug "Org Member #{org} -> #{x['login']} exists"
        end
      end

      persister.find(:org_members, {'org' => org}).map{|o| retrieve_org(o['login'])}
    end

    # Retrieve all comments for a single commit
    def retrieve_commit_comments(owner, repo, sha)
      retrieved_comments = paged_api_request(ghurl "repos/#{owner}/#{repo}/commits/#{sha}/comments")

      retrieved_comments.each { |x|
        if persister.find(:commit_comments, { 'commit_id' => x['commit_id'],
                                              'id' => x['id']}).empty?
          persister.store(:commit_comments, x)
        end
      }
      persister.find(:commit_comments, {'commit_id' => sha})
    end

    # Retrieve a single comment
    def retrieve_commit_comment(owner, repo, sha, id)

      comment = persister.find(:commit_comments, {'commit_id' => sha,
                                                  'id' => id}).first
      if comment.nil?
        r = api_request(ghurl "repos/#{owner}/#{repo}/comments/#{id}")

        if r.nil? or r.empty?
          warn "Could not find commit_comment #{id}. Deleted?"
          return
        end

        persister.store(:commit_comments, r)
        info "Added commit_comment #{r['commit_id']} -> #{r['id']}"
        persister.find(:commit_comments, {'commit_id' => sha, 'id' => id}).first
      else
        debug "Commit comment #{comment['commit_id']} -> #{comment['id']} exists"
        comment
      end
    end

    # Retrieve all watchers for a repository
    def retrieve_watchers(user, repo)
      repo_bound_items(user, repo, :watchers,
                       ["repos/#{user}/#{repo}/stargazers"],
                       {'repo' => repo, 'owner' => user},
                       'login', item = nil, refresh = false, order = :desc)
    end

    # Retrieve a single watcher for a repository
    def retrieve_watcher(user, repo, watcher)
      repo_bound_item(user, repo, watcher, :watchers,
                      ["repos/#{user}/#{repo}/stargazers"],
                      {'repo' => repo, 'owner' => user},
                      'login', order = :desc)
    end

    def retrieve_pull_requests(user, repo, refr = false)
      open = "repos/#{user}/#{repo}/pulls"
      closed = "repos/#{user}/#{repo}/pulls?state=closed"
      repo_bound_items(user, repo, :pull_requests,
                       [open, closed],
                       {'repo' => repo, 'owner' => user},
                       'number', item = nil, refresh = refr, order = :asc)
    end

    def retrieve_pull_request(user, repo, pullreq_id)
      open = "repos/#{user}/#{repo}/pulls"
      closed = "repos/#{user}/#{repo}/pulls?state=closed"
      repo_bound_item(user, repo, pullreq_id, :pull_requests,
                      [open, closed],
                      {'repo' => repo, 'owner' => user,
                       'number' => pullreq_id},
                      'number')
    end

    def retrieve_forks(user, repo)
      repo_bound_items(user, repo, :forks,
                       ["repos/#{user}/#{repo}/forks"],
                       {'repo' => repo, 'owner' => user},
                       'id', item = nil, refresh = false, order = :asc)
    end

    def retrieve_fork(user, repo, fork_id)
      repo_bound_item(user, repo, fork_id, :forks,
                       ["repos/#{user}/#{repo}/forks"],
                       {'repo' => repo, 'owner' => user},
                       'id')
    end

    def retrieve_pull_req_commits(user, repo, pullreq_id)
      pr_commits = paged_api_request(ghurl "repos/#{user}/#{repo}/pulls/#{pullreq_id}/commits")

      pr_commits.map do |x|
        head_user = x['url'].split(/\//)[4]
        head_repo = x['url'].split(/\//)[5]

        retrieve_commit(head_repo, x['sha'], head_user)
      end.select{|x| not x.nil?}
    end

    def retrieve_pull_req_comments(owner, repo, pullreq_id)
      review_comments_url = ghurl "repos/#{owner}/#{repo}/pulls/#{pullreq_id}/comments"

      url = review_comments_url
      retrieved_comments = paged_api_request url

      retrieved_comments.each do |x|
        x['owner'] = owner
        x['repo'] = repo
        x['pullreq_id'] = pullreq_id.to_i

        if persister.find(:pull_request_comments, {'owner' => owner,
                                                   'repo' => repo,
                                                   'pullreq_id' => pullreq_id,
                                                   'id' => x['id']}).empty?
          persister.store(:pull_request_comments, x)
        end
      end

      persister.find(:pull_request_comments, {'owner' => owner, 'repo' => repo,
                                              'pullreq_id' => pullreq_id})
    end

    def retrieve_pull_req_comment(owner, repo, pullreq_id, comment_id)
      comment = persister.find(:pull_request_comments, {'repo' => repo,
                                                 'owner' => owner,
                                                 'pullreq_id' => pullreq_id.to_i,
                                                 'id' => comment_id}).first
      if comment.nil?
        r = api_request(ghurl "repos/#{owner}/#{repo}/pulls/comments/#{comment_id}")

        if r.nil? or r.empty?
          warn "Could not find pullreq_comment #{owner}/#{repo} #{pullreq_id}->#{comment_id}. Deleted?"
          return
        end

        r['repo'] = repo
        r['owner'] = owner
        r['pullreq_id'] = pullreq_id.to_i
        persister.store(:pull_request_comments, r)
        info "Added pullreq_comment #{owner}/#{repo} #{pullreq_id}->#{comment_id}"
        persister.find(:pull_request_comments, {'repo' => repo, 'owner' => owner,
                                         'pullreq_id' => pullreq_id.to_i,
                                         'id' => comment_id}).first
      else
        debug "Pull request comment #{owner}/#{repo} #{pullreq_id}->#{comment_id} exists"
        comment
      end
    end

    def retrieve_issues(user, repo, refr = false)
      open = "repos/#{user}/#{repo}/issues"
      closed = "repos/#{user}/#{repo}/issues?state=closed"
      repo_bound_items(user, repo, :issues,
                       [open, closed],
                       {'repo' => repo, 'owner' => user},
                       'number', item = nil, refresh = refr, order = :asc)
    end

    def retrieve_issue(user, repo, issue_id)
      open = "repos/#{user}/#{repo}/issues"
      closed = "repos/#{user}/#{repo}/issues?state=closed"
      repo_bound_item(user, repo, issue_id, :issues,
                      [open, closed],
                      {'repo' => repo, 'owner' => user},
                      'number')
    end

    def retrieve_issue_events(owner, repo, issue_id)
      url = ghurl "repos/#{owner}/#{repo}/issues/#{issue_id}/events"
      retrieved_events = paged_api_request url

      issue_events = retrieved_events.map { |x|
        x['owner'] = owner
        x['repo'] = repo
        x['issue_id'] = issue_id

        if persister.find(:issue_events, {'owner' => owner,
                                          'repo' => repo,
                                          'issue_id' => issue_id,
                                          'id' => x['id']}).empty?
          info "Added issue_event #{owner}/#{repo} #{issue_id}->#{x['id']}"
          persister.store(:issue_events, x)
        end
        x
      }
      a = persister.find(:issue_events, {'owner' => owner, 'repo' => repo,
                                         'issue_id' => issue_id})
      if a.empty? then issue_events else a end
    end

    def retrieve_issue_event(owner, repo, issue_id, event_id)
      event = persister.find(:issue_events, {'repo' => repo,
                                             'owner' => owner,
                                             'issue_id' => issue_id,
                                             'id' => event_id}).first
      if event.nil?
        r = api_request(ghurl "repos/#{owner}/#{repo}/issues/events/#{event_id}")

        if r.nil? or r.empty?
          warn "Could not find issue_event #{owner}/#{repo} #{issue_id}->#{event_id}. Deleted?"
          return
        end

        r['repo'] = repo
        r['owner'] = owner
        r['issue_id'] = issue_id
        persister.store(:issue_events, r)
        info "Added issue_event #{owner}/#{repo} #{issue_id}->#{event_id}"
        a = persister.find(:issue_events, {'repo' => repo, 'owner' => owner,
                                       'issue_id' => issue_id,
                                       'id' => event_id}).first
        if a.nil? then r else a end
      else
        debug "Issue event #{owner}/#{repo} #{issue_id}->#{event_id} exists"
        event
      end
    end

    def retrieve_issue_comments(owner, repo, issue_id)
      url = ghurl "repos/#{owner}/#{repo}/issues/#{issue_id}/comments"
      retrieved_comments = paged_api_request url

      comments = retrieved_comments.each { |x|
        x['owner'] = owner
        x['repo'] = repo
        x['issue_id'] = issue_id

        if persister.find(:issue_comments, {'owner' => owner,
                                            'repo' => repo,
                                            'issue_id' => issue_id,
                                            'id' => x['id']}).empty?
          persister.store(:issue_comments, x)
        end
        x
      }
      a = persister.find(:issue_comments, {'owner' => owner, 'repo' => repo,
                                           'issue_id' => issue_id})
      if a.empty? then comments else a end
    end

    def retrieve_issue_comment(owner, repo, issue_id, comment_id)
      comment = persister.find(:issue_comments, {'repo' => repo,
                                                 'owner' => owner,
                                                 'issue_id' => issue_id,
                                                 'id' => comment_id}).first
      if comment.nil?
        r = api_request(ghurl "repos/#{owner}/#{repo}/issues/comments/#{comment_id}")

        if r.nil? or r.empty?
          warn "Could not find issue_comment #{owner}/#{repo} #{issue_id}->#{comment_id}. Deleted?"
          return
        end

        r['repo'] = repo
        r['owner'] = owner
        r['issue_id'] = issue_id
        persister.store(:issue_comments, r)
        info "Added issue_comment #{owner}/#{repo} #{issue_id}->#{comment_id}"
        a = persister.find(:issue_comments, {'repo' => repo, 'owner' => owner,
                                         'issue_id' => issue_id,
                                         'id' => comment_id}).first
        if a.nil? then r else a end
      else
        debug "Issue comment #{owner}/#{repo} #{issue_id}->#{comment_id} exists"
        comment
      end
    end

    def retrieve_repo_labels(owner, repo, refr = false)
      repo_bound_items(owner, repo, :repo_labels,
                       ["repos/#{owner}/#{repo}/labels"],
                       {'repo' => repo, 'owner' => owner},
                       'name', item = nil, refresh = refr, order = :asc)
    end

    def retrieve_repo_label(owner, repo, name)
      repo_bound_item(owner, repo, name, :repo_labels,
                       ["repos/#{owner}/#{repo}/labels"],
                       {'repo' => repo, 'owner' => owner},
                       'name')
    end

    def retrieve_issue_labels(owner, repo, issue_id)
      url = ghurl("repos/#{owner}/#{repo}/issues/#{issue_id}/labels")
      paged_api_request(url)
    end

    def retrieve_topics(owner, repo)
      # volatile: currently available with api preview
      # https://developer.github.com/v3/repos/#list-all-topics-for-a-repository
      stored_topics = persister.find(:topics, {'owner' => owner, 'repo' => repo })

      url = ghurl("repos/#{owner}/#{repo}/topics")
      r = api_request(url, media_type = "application/vnd.github.mercy-preview+json")

      if r.nil? or r.empty? or r['names'].nil? or r['names'].empty?
        warn "No topics for #{owner}/#{repo}"
        return []
      end

      topics = r['names']
      if topics.nil? or topics.empty?
        return []
      end

      topics.each do |topic|
        if stored_topics.select{|x| x['topic'] == topic}.empty?
          topic_entry = {
            :owner => owner,
            :repo  => repo,
            :topic => topic
          }
          persister.store(:topics, topic_entry)
          info "Added topic #{topic} -> #{owner}/#{repo}"
        else
          debug "Topic #{topic} -> #{owner}/#{repo} exists"
        end
      end

      r['names']
    end

    # Get current Github events
    def get_events
      api_request "https://api.github.com/events"
    end

    # Get all events for the specified repo.
    # GitHub will only return 90 days of events
    def get_repo_events(owner, repo)
      url = ghurl("repos/#{owner}/#{repo}/events")
      r = paged_api_request(url)

      r.each do |e|
        unless get_event(e['id']).empty?
          debug "Repository event #{owner}/#{repo} -> #{e['type']}-#{e['id']} already exists"
        else
          persister.store(:events, e)
          info "Added event for repository #{owner}/#{repo} -> #{e['type']}-#{e['id']}"
        end
      end

      persister.find(:events, {'repo.name' => "#{owner}/#{repo}"})

    end

    # Get a specific event by +id+.
    def get_event(id)
      persister.find(:events, {'id' => id})
    end

    # Retrieve diff between two branches. If either branch name is not provided
    # the branch name is resolved to the corresponding default branch
    def retrieve_master_branch_diff(owner, repo, branch, parent_owner, parent_repo, parent_branch)
      branch   = retrieve_default_branch(owner, repo) if branch.nil?
      parent_branch = retrieve_default_branch(parent_owner, parent_repo) if parent_branch.nil?
      return nil if branch.nil? or parent_branch.nil?

      cmp_url = "https://api.github.com/repos/#{parent_owner}/#{parent_repo}/compare/#{parent_branch}...#{owner}:#{branch}"
      api_request(cmp_url)
    end

    # Retrieve the default branch for a repo. If nothing is retrieved, 'master' is returned
    def retrieve_default_branch(owner, repo, refresh = false)
      retrieved = retrieve_repo(owner, repo, refresh)
      return nil if retrieved.nil?

      master_branch = 'master'
      if retrieved['default_branch'].nil?
        # The currently stored repo entry has been created before the
        # default_branch field was added to the schema
        retrieved = retrieve_repo(owner, repo, true)
        return nil if retrieved.nil?
      end
      master_branch = retrieved['default_branch'] unless retrieved.nil?
      master_branch
    end

    private

    def restricted_page_request(url, pages)
      if pages != -1
        paged_api_request(url, pages)
      else
        paged_api_request(url)
      end
    end

    def repo_bound_items(user, repo, entity, urls, selector, discriminator,
        item_id = nil, refresh = false, order = :asc, media_type = '')

       urls.each do |url|
        total_pages = num_pages(ghurl url)

        page_range = if order == :asc
                       (1..total_pages)
                     else
                       total_pages.downto(1)
                     end

        page_range.each do |page|
          items = api_request(ghurl(url, page), media_type)
          break if items.nil?

          items.each do |x|
            x['repo'] = repo
            x['owner'] = user

            instances = repo_bound_instance(entity, selector,
                                            discriminator, x[discriminator])
            exists = !instances.empty?

            unless exists
              x = api_request(x['url'], media_type)
              break if x.nil?
              x['repo'] = repo
              x['owner'] = user
              persister.store(entity, x)
              info "Added #{entity} #{user}/#{repo} -> #{x[discriminator]}"
            else
              if refresh
                instances.each do |i|

                  id = if i[discriminator].to_i.to_s != i[discriminator]
                         i[discriminator] # item_id is int
                       else
                         i[discriminator].to_i # convert to int
                       end

                  instance_selector = selector.merge({discriminator => id})
                  persister.upsert(entity, instance_selector, x)
                  debug "Refreshing #{entity} #{user}/#{repo} -> #{x[discriminator]}"
                end
              else
                debug "#{entity} #{user}/#{repo} -> #{x[discriminator]} exists"
              end
            end

            # If we are just looking for a single item, give the method a chance
            # to return as soon as we find it. This is to avoid loading all
            # items before we actually search for what we are looking for.
            unless item_id.nil?
              a = repo_bound_instance(entity, selector, discriminator, item_id)
              unless a.empty?
                return a
              end
            end
          end
        end
      end

      if item_id.nil?
        persister.find(entity, selector)
      else
        # If the item we are looking for has been found, the method should
        # have returned earlier. So just return an empty result to indicate
        # that the item has not been found.
        []
      end
    end

    def repo_bound_item(user, repo, item_id, entity, url, selector,
                        discriminator, order = :asc, media_type = '')
      stored_item = repo_bound_instance(entity, selector, discriminator, item_id)

      r = if stored_item.empty?
            repo_bound_items(user, repo, entity, url, selector, discriminator,
                             item_id, false, order, media_type).first
          else
            stored_item.first
          end
      if r.nil?
        warn "Could not find #{entity} #{user}/#{repo} -> #{item_id}. Deleted?"
      end
      r
    end

    def repo_bound_instance(entity, selector, discriminator, item_id)

      id = if item_id.to_i.to_s != item_id
             item_id # item_id is int
           else
             item_id.to_i # convert to int
           end

      instance_selector = selector.merge({discriminator => id})
      result = persister.find(entity, instance_selector)
      if result.empty?
        # Try without type conversions. Useful when the discriminator type
        # is string and an item_id that can be converted to int is passed.
        # Having no types sucks occasionaly...
        instance_selector = selector.merge({discriminator => item_id})
        persister.find(entity, instance_selector)
      else
        result
      end
    end

    def ghurl(path, page = -1, per_page = 100)
      if page > 0
        if path.include?('?')
          path = path + "&page=#{page}&per_page=#{per_page}"
        else
          path = path + "?page=#{page}&per_page=#{per_page}"
        end
        config(:mirror_urlbase) + path
      else
        if path.include?('?')
          path = path + "&per_page=#{per_page}"
        else
          path = path + "?per_page=#{per_page}"
        end
        config(:mirror_urlbase) + path
      end
    end

  end
end


require 'rubygems'
require 'trollop'
require 'bunny'
require 'etc'

require 'ghtorrent/settings'
require 'version'

module GHTorrent

  # Base class for all GHTorrent command line utilities. Provides basic command
  # line argument parsing and command bootstraping support. The order of
  # initialization is the following:
  # prepare_options
  # validate
  # go
  class Command

    include GHTorrent::Settings
    include GHTorrent::Logging

    # Specify the run method for subclasses.
    class << self
      def run(args = ARGV)
        attr_accessor :args
        attr_accessor :settings
        attr_accessor :name
        attr_accessor :options

        command = new()

        command.name = self.class.name
        command.args = args

        command.process_options
        command.validate

        command.settings = YAML::load_file command.options[:config]

        unless command.options[:addr].nil?
          command.settings = command.override_config(command.settings,
                                                     :attach_ip,
                                                     command.options[:addr])
        end

        unless command.options[:token].nil?
          command.settings = command.override_config(command.settings,
                                                     :github_token,
                                                     command.options[:token])
        end

        unless command.options[:req_limit].nil?
          command.settings = command.override_config(command.settings,
                                                     :req_limit,
                                                     command.options[:req_limit])
        end

        unless command.options[:uniq].nil?
          command.settings = command.override_config(command.settings,
                                                     :logging_uniq,
                                                     command.options[:uniq])
        end

        begin
          command.go
        rescue => e
          STDERR.puts e.message
          if command.options.verbose
            STDERR.puts e.backtrace.join("\n")
          else
            STDERR.puts e.backtrace[0]
          end
          exit 1
        end
      end
    end

    # Specify and parse top-level command line options.
    def process_options
      command = self
      @options = Trollop::options(command.args) do

        command.prepare_options(self)

        banner <<-END
Standard options:
        END

        opt :config, 'config.yaml file location', :short => 'c',
            :default => 'config.yaml'
        opt :verbose, 'verbose mode', :short => 'v'
        opt :addr, 'IP address to use for performing requests', :short => 'a',
            :type => String
        opt :token, 'GitHub OAuth token',
            :type => String, :short => 't'
        opt :req_limit, 'Number or requests to leave on any provided account (in reqs/hour)',
            :type => Integer, :short => 'l'
        opt :uniq, 'Unique name for this command. Will appear in logs.',
            :type => String, :short => 'u'
      end
    end

    # Get the version of the project
    def version
      IO.read(File.join(File.dirname(__FILE__), '..', '..', 'VERSION'))
    end

    # This method should be overriden by subclasses in order to specify,
    # using trollop, the supported command line options
    def prepare_options(options)
    end

    # Examine the validity of the provided options in the context of the
    # executed command. Subclasses can also call super to also invoke the checks
    # provided by this class.
    def validate
      if options[:config].nil?
        unless (File.exist?("config.yaml"))
          Trollop::die "No config file in default location (#{Dir.pwd}). You
                        need to specify the #{:config} parameter. Read the
                        documentation on how to create a config.yaml file."
        end
      else
        Trollop::die "Cannot find file #{options[:config]}" \
          unless File.exist?(options[:config])
      end

      unless @options[:user].nil?
        if not Process.uid == 0
          Trollop::die "Option --user (-u) can only be specified by root"
        end
          begin
            Etc.getpwnam(@options[:user])
          rescue ArgumentError
            Trollop::die "No such user: #{@options[:user]}"
          end
      end
    end

    # Name of the command that is currently being executed.
    def command_name
      File.basename($0)
    end

    # The actual command code.
    def go
    end

    # Specify a handler to incoming messages from a connection to a queue.
    #
    # @param queue [String] the queue name to bind to
    # @param key [String] routing key for msgs for binding the queue to the exchange.
    # @param ack [Symbol] when should acks be send, :before or :after the block returns
    # @param block [Block]: A block accepting one argument (the message)
    def queue_client(queue, key = queue, ack = :after, block)

      stopped = false
      while not stopped
        begin
          conn = Bunny.new(:host => config(:amqp_host),
                           :port => config(:amqp_port),
                           :username => config(:amqp_username),
                           :password => config(:amqp_password))
          conn.start

          ch  = conn.create_channel
          debug "Queue setting prefetch to #{config(:amqp_prefetch)}"
          ch.prefetch(config(:amqp_prefetch))
          debug "Queue connection to #{config(:amqp_host)} succeeded"

          x = ch.topic(config(:amqp_exchange), :durable => true,
                       :auto_delete => false)
          q = ch.queue(queue, :durable => true)
          q.bind(x, :routing_key => key)

          q.subscribe(:block => true,
                      :manual_ack => true) do |delivery_info, properties, msg|

            if ack == :before
              ch.acknowledge(delivery_info.delivery_tag)
            end

            begin
              block.call(msg)
            ensure
              if ack != :before
                ch.acknowledge(delivery_info.delivery_tag)
              end
            end
          end

        rescue Bunny::TCPConnectionFailed => e
          warn "Connection to #{config(:amqp_host)} failed. Retrying in 1 sec"
          sleep(1)
        rescue Bunny::PossibleAuthenticationFailureError => e
          warn "Could not authenticate as #{conn.username}"
        rescue Bunny::NotFound, Bunny::AccessRefused, Bunny::PreconditionFailed => e
          warn "Channel error: #{e}. Retrying in 1 sec"
          sleep(1)
        rescue Interrupt => _
          stopped = true
        rescue StandardError => e
          raise e
        end
      end

      ch.close unless ch.nil?
      conn.close unless conn.nil?

    end

    def override_config(config_file, setting, new_value)
      puts "Overriding configuration #{setting}=#{config(setting)} with new value #{new_value}"
      super(config_file, setting, new_value)
    end

  end

end
# vim: set sta sts=2 shiftwidth=2 sw=2 et ai :


require 'ghtorrent/hash'

module GHTorrent
  module Utils

    def self.included(other)
      other.extend self
    end

    # Read the value for a key whose format is "foo.bar.baz" from a hierarchical
    # map, where a dot represents one level deep in the hierarchy.
    def read_value(from, key)
      return from if key.nil? or key == ""

      key.split(/\./).reduce({}) do |acc, x|
        unless acc.nil?
          if acc.empty?
            # Initial run
            acc = from[x]
          else
            if acc.has_key?(x)
              acc = acc[x]
            else
              # Some intermediate key does not exist
              return nil
            end
          end
        else
          # Some intermediate key returned a null value
          # This indicates a malformed entry
          return nil
        end
      end
    end

    # Overwrite an existing +key+ whose format is "foo.bar" (where a dot
    # represents one level deep in the hierarchy) in hash +to+ with +value+.
    # If the key does not exist, it will be added at the appropriate depth level
    def write_value(to, key, value)
      return to if key.nil? or key == ""

      prev = nil
      key.split(/\./).reverse.each {|x|
        a = Hash.new
        a[x] = if prev.nil? then value else prev end
        prev = a
        a
      }

      to.merge_recursive(prev)
    end

    def user_type(type)
      if type == "User"
        "USR"
      else
        "ORG"
      end
    end
  end
end


require 'open-uri'
require 'json'
require 'uri'

module GHTorrent
  module Geolocator
    module OSM

      def format_url(location)
        URI.escape("http://nominatim.openstreetmap.org/search/#{location}?format=json&addressdetails=1&accept_language=en")
      end

      def parse_geolocation_result(location, geocoded)
        if geocoded.empty?
          debug "Geolocation request #{location} returned successfully but no location was found"
          geo       = EMPTY_LOCATION
          geo[:key] = location
          return geo
        else
          debug "Geolocation request for #{location} returned #{geocoded.size} places"
        end

        geocoded = geocoded.sort { |x, y| y['importance'] <=> x['importance'] }.first

        geo = {
            :key          => location,
            :long         => geocoded['lon'],
            :lat          => geocoded['lat'],
            :city         => geocoded['address']['city'],
            :country      => geocoded['address']['country'],
            :state        => geocoded['address']['state'],
            :country_code => geocoded['address']['country_code'],
            :status       => :ok
        }
        geo

      end
    end
  end
end

module GHTorrent
  module Geolocator
    module Bing

      COUNTRY_CODES = File.open(File.join(File.dirname(__FILE__), 'country_codes.txt'), "r:UTF-8").\
        readlines.\
        inject({}) do |acc, x|
          name, code = x.split(/','/)
          acc.merge(name.tr("'", '') => code.strip.tr("'", ''))
        end

    CONFIDENCE_ORDER = {'High' => 3, 'Medium' => 2, 'Low' => 1}

      def format_url(location)
        URI.escape("http://dev.virtualearth.net/REST/v1/Locations?q=#{location}&key=#{config(:geolocation_bing_key)}")
      end

      def parse_geolocation_result(location, geocoded)
        if geocoded['resourceSets'].nil? or geocoded['resourceSets'][0].nil? or
            geocoded['resourceSets'][0]['estimatedTotal'] < 1
          debug "Geolocation request for #{location} returned successfully but no location was found"
          geo       = EMPTY_LOCATION
          geo[:key] = location
          return geo
        else
          debug "Geolocation request for #{location} returned #{geocoded['resourceSets'][0]['estimatedTotal']} places"
        end

        details = geocoded['resourceSets'][0]['resources'].map do |x|
          x['confidence'] = CONFIDENCE_ORDER[x['confidence']]
          x
        end.sort_by do |x|
          x['confidence']
        end.select do |x|
          x['confidence'] == 3
        end.first

        geo = {
            :key          => location,
            :long         => details['point']['coordinates'][1],
            :lat          => details['point']['coordinates'][0],
            :city         => details['address']['locality'],
            :country      => details['address']['countryRegion'],
            :state        => details['address']['adminDistrict'],
            :country_code => COUNTRY_CODES[details['address']['countryRegion']],
            :status       => :ok
        }
        geo
      end
    end
  end
end

module GHTorrent
  module Geolocator
    module GMaps

      def format_url(location)
        URI.escape("https://maps.googleapis.com/maps/api/geocode/json?key=#{config(:geolocation_gmaps_key)}&address=#{location}")
      end

      def parse_geolocation_result(location, geocoded)
        details = geocoded['results'].first
        city_hash = details['address_components'].find{|x| x['types'].include? 'locality'}
        country_hash = details['address_components'].find{|x| x['types'].include? 'country'}
        admin_area_hash = details['address_components'].select{|x| not x['types'].grep(/administrative_area/).empty?}.first

        geo = {
            :key          => location,
            :long         => details['geometry']['location']['lng'],
            :lat          => details['geometry']['location']['lat'],
            :city         => unless city_hash.nil? then city_hash['long_name'] end,
            :country      => unless country_hash.nil? then country_hash['long_name'] end,
            :state        => unless admin_area_hash.nil? then admin_area_hash['long_name'] end,
            :country_code => unless country_hash.nil? then country_hash['short_name'].downcase end,
            :status       => :ok
        }

        geo
      end
    end
  end
end

module GHTorrent
  module Geolocator

    include GHTorrent::Settings

    EMPTY_LOCATION = {
        :key          => nil,
        :long         => nil,
        :lat          => nil,
        :city         => nil,
        :country      => nil,
        :state        => nil,
        :country_code => nil,
        :status       => :failed
    }

    def persister
      raise Exception.new("Unimplemented")
    end

    # Given a location string, it returns a hash like the following:
    #
    # input: Larisa, Greece
    # output:
    # {
    #   :key          => "larisa",
    #   :long         => 22.41629981994629,
    #   :lat          => 39.61040115356445,
    #   :city         => "Larisa",
    #   :country      => "Greece",
    #   :state        => "Thessaly",
    #   :country_code => "gr",
    #   :status       => :ok
    # }
    # Uses aggressive caching
    def geolocate(location: nil, wait: config(:geolocation_wait).to_i, from_cache: true)
      return EMPTY_LOCATION if location.nil? or location == ''
      location = location_filter(location)

      geo = []
      if from_cache
        geo = persister.find(:geo_cache, {'key' => location})
      end

      if geo.empty?

        if config(:geolocation_service) == 'gmaps'
          self.class.send :include, GHTorrent::Geolocator::GMaps
        elsif config(:geolocation_service) == 'bing'
          self.class.send :include, GHTorrent::Geolocator::Bing
        else
          self.class.send :include, GHTorrent::Geolocator::OSM
        end

        begin
          ts  = Time.now
          url = format_url(location)
          req = open(url)
          p   = JSON.parse(req.read)
          geo = parse_geolocation_result(location, p)

          info "Successful geolocation request. Location: #{location}"
        rescue StandardError => e
          warn "Failed geolocation request. Location: #{location}"
          geo       = EMPTY_LOCATION
          geo[:key] = location
        ensure
          in_db_geo = persister.find(:geo_cache, {'key' => location}).first

          if in_db_geo.nil?
            begin

              geo[:updated_at] = Time.now

              persister.store(:geo_cache, geo)
            rescue StandardError => e
              warn "Could not save location #{location} -> #{geo}: #{e.message}"
            end
          end

          info "Added location key '#{location}' -> #{geo[:status]}"
          taken    = Time.now.to_f - ts.to_f
          to_sleep = wait - taken
          sleep(to_sleep) if to_sleep > 0
        end
      else
        geo = geo[0]
        debug "Location with key '#{location}' exists"
      end

      geo
    end

    # Standard filtering on all locations used by GHTorrent
    def location_filter(location)
      return nil if location.nil?
      location.\
        strip.\
        downcase.\
        tr('#"<>[]', '').\
        gsub(/^[0-9,\/().:]*/, '').\
        gsub(/ +/, ' ').\
        gsub(/,([a-z]*)/, '\1')
    end
  end
end


# frozen_string_literal: true

module Karafka
  # Base responder from which all Karafka responders should inherit
  # Similar to Rails responders concept. It allows us to design flow from one app to another
  # by isolating what responses should be sent (and where) based on a given action
  # It differs from Rails responders in the way it works: in std http request we can have one
  # response, here we can have unlimited number of them
  #
  # It has a simple API for defining where should we respond (and if it is required)
  #
  # @example Basic usage (each registered topic is required to be used by default)
  #   class Responder < BaseResponder
  #     topic :new_action
  #
  #     def respond(data)
  #       respond_to :new_action, data
  #     end
  #   end
  #
  # @example Responding to a topic with extra options
  #   class Responder < BaseResponder
  #     topic :new_action
  #
  #     def respond(data)
  #       respond_to :new_action, data, partition_key: 'thing'
  #     end
  #   end
  #
  # @example Marking topic as not required (we won't have to use it)
  #   class Responder < BaseResponder
  #     topic :required_topic
  #     topic :new_action, required: false
  #
  #     def respond(data)
  #       respond_to :required_topic, data
  #     end
  #   end
  #
  # @example Multiple times used topic
  #   class Responder < BaseResponder
  #     topic :required_topic
  #
  #     def respond(data)
  #       data.each do |subset|
  #         respond_to :required_topic, subset
  #       end
  #     end
  #   end
  #
  # @example Specify serializer for a topic
  #   class Responder < BaseResponder
  #     topic :xml_topic, serializer: MyXMLSerializer
  #
  #     def respond(data)
  #       data.each do |subset|
  #         respond_to :xml_topic, subset
  #       end
  #     end
  #   end
  #
  # @example Accept multiple arguments to a respond method
  #   class Responder < BaseResponder
  #     topic :users_actions
  #     topic :articles_viewed
  #
  #     def respond(user, article)
  #       respond_to :users_actions, user
  #       respond_to :articles_viewed, article
  #     end
  #   end
  class BaseResponder
    class << self
      # Definitions of all topics that we want to be able to use in this responder should go here
      attr_accessor :topics
      # Schema that we can use to control and/or require some additional details upon options
      # that are being passed to the producer. This can be in particular useful if we want to make
      # sure that for example partition_key is always present.
      attr_accessor :options_schema

      # Registers a topic as on to which we will be able to respond
      # @param topic_name [Symbol, String] name of topic to which we want to respond
      # @param options [Hash] hash with optional configuration details
      def topic(topic_name, options = {})
        options[:serializer] ||= Karafka::App.config.serializer
        options[:registered] = true
        self.topics ||= {}
        topic_obj = Responders::Topic.new(topic_name, options)
        self.topics[topic_obj.name] = topic_obj
      end

      # A simple alias for easier standalone responder usage.
      # Instead of building it with new.call it allows (in case of using JSON serializer)
      # to just run it directly from the class level
      # @param data Anything that we want to respond with
      # @example Send user data with a responder
      #   UsersCreatedResponder.call(@created_user)
      def call(*data)
        # Just in case there were no topics defined for a responder, we initialize with
        # empty hash not to handle a nil case
        self.topics ||= {}
        new.call(*data)
      end
    end

    attr_reader :messages_buffer

    # Creates a responder object
    # @return [Karafka::BaseResponder] base responder descendant responder
    def initialize
      @messages_buffer = {}
    end

    # Performs respond and validates that all the response requirement were met
    # @param data Anything that we want to respond with
    # @note We know that validators should be executed also before sending data to topics, however
    #   the implementation gets way more complicated then, that's why we check after everything
    #   was sent using responder
    # @example Send user data with a responder
    #   UsersCreatedResponder.new.call(@created_user)
    # @example Send user data with a responder using non default Parser
    #   UsersCreatedResponder.new(MyParser).call(@created_user)
    def call(*data)
      respond(*data)
      validate_usage!
      validate_options!
      deliver!
    end

    private

    # Checks if we met all the topics requirements. It will fail if we didn't send a message to
    # a registered required topic, etc.
    def validate_usage!
      registered_topics = self.class.topics.map do |name, topic|
        topic.to_h.merge!(
          usage_count: messages_buffer[name]&.count || 0
        )
      end

      used_topics = messages_buffer.map do |name, usage|
        topic = self.class.topics[name] || Responders::Topic.new(name, registered: false)
        topic.to_h.merge!(usage_count: usage.count)
      end

      result = Karafka::Schemas::ResponderUsage.call(
        registered_topics: registered_topics,
        used_topics: used_topics
      )

      return if result.success?

      raise Karafka::Errors::InvalidResponderUsageError, result.errors
    end

    # Checks if we met all the options requirements before sending them to the producer.
    def validate_options!
      return true unless self.class.options_schema

      messages_buffer.each_value do |messages_set|
        messages_set.each do |message_data|
          result = self.class.options_schema.call(message_data.last)
          next if result.success?

          raise Karafka::Errors::InvalidResponderMessageOptionsError, result.errors
        end
      end
    end

    # Takes all the messages from the buffer and delivers them one by one
    # @note This method is executed after the validation, so we're sure that
    #   what we send is legit and it will go to a proper topics
    def deliver!
      messages_buffer.each_value do |data_elements|
        data_elements.each do |data, options|
          # We map this topic name, so it will match namespaced/etc topic in Kafka
          # @note By default will not change topic (if default mapper used)
          mapped_topic = Karafka::App.config.topic_mapper.outgoing(options[:topic])
          external_options = options.merge(topic: mapped_topic)
          producer(options).call(data, external_options)
        end
      end
    end

    # Method that needs to be implemented in a subclass. It should handle responding
    #   on registered topics
    # @raise [NotImplementedError] This method needs to be implemented in a subclass
    def respond(*_data)
      raise NotImplementedError, 'Implement this in a subclass'
    end

    # This method allow us to respond to a single topic with a given data. It can be used
    # as many times as we need. Especially when we have 1:n flow
    # @param topic [Symbol, String] topic to which we want to respond
    # @param data [String, Object] string or object that we want to send
    # @param options [Hash] options for waterdrop (e.g. partition_key).
    # @note Respond to does not accept multiple data arguments.
    def respond_to(topic, data, options = {})
      # We normalize the format to string, as WaterDrop and Ruby-Kafka support only
      # string topics
      topic = topic.to_s

      messages_buffer[topic] ||= []
      messages_buffer[topic] << [
        self.class.topics[topic].serializer.call(data),
        options.merge(topic: topic)
      ]
    end

    # @param options [Hash] options for waterdrop
    # @return [Class] WaterDrop producer (sync or async based on the settings)
    def producer(options)
      if self.class.topics[options[:topic]].async?
        WaterDrop::AsyncProducer
      else
        WaterDrop::SyncProducer
      end
    end
  end
end


# frozen_string_literal: true

module Karafka
  # Class used to catch signals from ruby Signal class in order to manage Karafka stop
  # @note There might be only one process - this class is a singleton
  class Process
    include Singleton

    # Signal types that we handle
    HANDLED_SIGNALS = %i[
      SIGINT
      SIGQUIT
      SIGTERM
    ].freeze

    HANDLED_SIGNALS.each do |signal|
      # Assigns a callback that will happen when certain signal will be send
      # to Karafka server instance
      # @note It does not define the callback itself -it needs to be passed in a block
      # @example Define an action that should be taken on_sigint
      #   process.on_sigint do
      #     Karafka.logger.info('Log something here')
      #     exit
      #   end
      define_method :"on_#{signal.to_s.downcase}" do |&block|
        @callbacks[signal] << block
      end
    end

    # Creates an instance of process and creates empty hash for callbacks
    def initialize
      @callbacks = Hash.new { |hsh, key| hsh[key] = [] }
    end

    # Method catches all HANDLED_SIGNALS and performs appropriate callbacks (if defined)
    # @note If there are no callbacks, this method will just ignore a given signal that was sent
    def supervise
      HANDLED_SIGNALS.each { |signal| trap_signal(signal) }
    end

    private

    # Traps a single signal and performs callbacks (if any) or just ignores this signal
    # @param [Symbol] signal type that we want to catch
    def trap_signal(signal)
      trap(signal) do
        notice_signal(signal)
        (@callbacks[signal] || []).each(&:call)
      end
    end

    # Informs monitoring about trapped signal
    # @param [Symbol] signal type that we received
    # @note We cannot perform logging from trap context, that's why
    #   we have to spin up a new thread to do this
    def notice_signal(signal)
      Thread.new do
        Karafka.monitor.instrument('process.notice_signal', caller: self, signal: signal)
      end
    end
  end
end


require 'tensor_stream/utils/py_ports'
##
# ruby port of https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/embedding_ops.py
#
module TensorStream
  module EmbeddingLookup
    include TensorStream::PyPorts

    ##
    # Looks up `ids` in a list of embedding tensors.
    def embedding_lookup(params, ids, partition_strategy: "mod", name: nil, validate_indices: true, max_norm: nil)
      _embedding_lookup_and_transform(params, ids, partition_strategy: partition_strategy, name: name, max_norm: max_norm, transform_fn: nil)
    end

    ##
    # Helper function for embedding_lookup and _compute_sampled_logits.
    def _embedding_lookup_and_transform(params, ids, partition_strategy: "mod", name: nil, max_norm: nil, transform_fn: nil)
      raise TensorStream::ValueError, "Need at least one param" if params.nil?

      params = [params] unless params.is_a?(Array)

      TensorStream.name_scope(name, "embedding_lookup", values: params + [ids]) do |name|
        np = params.size
        ids = TensorStream.convert_to_tensor(ids, name: "ids")
        if (np == 1) && (transform_fn.nil? || (ids.shape.size == 1))
          result = nil
          TensorStream.colocate_with(params[0]) do
            result = _clip(TensorStream.gather(params[0], ids, name: name), ids, max_norm)
            result = transform_fn.call(result) if transform_fn
          end

          return TensorStream.identity(result)
        else
          flat_ids = TensorStream.reshape(ids, [-1])
          original_indices = TensorStream.range(TensorStream.size(flat_ids))

          p_assignments = nil
          new_ids = nil

          if partition_strategy == "mod"
            p_assignments = flat_ids % np
            new_ids = floor_div(flat_ids, np)
          elsif partition_strategy == "div"
            raise "not yet supported!"
          else
            raise TensorStream::ValueError, "Unrecognized partition strategy: " + partition_strategy
          end

          p_assignments = TensorStream.cast(p_assignments, :int32)
          gather_ids = TensorStream.dynamic_partition(new_ids, p_assignments, np)
          pindices = TensorStream.dynamic_partition(original_indices, p_assignments, np)
          partitioned_result = []
          (0...np).each do |p|
            pids = gather_ids[p]
            result = nil
            TensorStream.colocate_with(params[p]) do
              result = TensorStream.gather(params[p], pids)
              if transform_fn
                # If transform_fn is provided, the clip_by_norm precedes
                # the transform and hence must be co-located. See below
                # for the counterpart if transform_fn is not proveded.
                result = transform_fn.call(_clip(result, pids, max_norm))
              end
            end
            partitioned_result << result
          end
          ret = TensorStream.dynamic_stitch(pindices, partitioned_result, name: name)

          if transform_fn.nil?
            element_shape_s = params[0].shape[1..-1]
            params[1..-1].each { |p| element_shape_s = element_shape_s.merge_with(p.shape[1..-1]) }
          else
            element_shape_s = ret.shape[1..-1]
          end

           # Compute the dynamic element shape.
          element_shape_d = if element_shape_s.fully_defined?
                               element_shape_s
                            elsif transform_fn.nil?
                              # It's important that we compute params[0].shape on the right device
                              # to avoid data motion.
                              TensorStream.colocate_with(params[0]) do
                                params_shape = TensorStream.shape(params[0])
                                params_shape[1..-1]
                              end
                            else
                              TensorStream.shape(ret)[1..-1]
                            end
          ret = TensorStream.reshape(ret, TensorStream.concat([TensorStream.shape(ids), element_shape_d], 0))
          ret = _clip(ret, ids, max_norm) unless transform_fn
          ret
        end
      end
    end

    def _clip(params, ids, max_norm)
      return params if max_norm.nil?

      ids_rank, ids_static = _rank(ids)
      params_rank, params_static = _rank(params)

      TensorStream.clip_by_norm(params, max_norm, axes: ids_static && params_static ? (ids_rank...params_rank).to_a : TensorStream.range(ids_rank, params_rank))
    end

    def _rank(x)
      rank = TensorStream.convert_to_tensor(x).shape.ndims
      if rank
        [rank, false]
      else
        [TensorStream.rank(x), false]
      end
    end
  end
end


require "yaml"

module TensorStream
  # A .pb graph deserializer
  class Protobuf
    def initialize
    end

    def load_from_string(buffer)
      evaluate_lines(buffer.split("\n").map(&:strip))
    end

    ##
    # parsers a protobuf file and spits out
    # a ruby hash
    def load(pbfile)
      f = File.new(pbfile, "r")
      lines = []
      while !f.eof? && (str = f.readline.strip)
        lines << str
      end
      evaluate_lines(lines)
    end

    def parse_value(value_node)
      return unless value_node["tensor"]

      evaluate_tensor_node(value_node["tensor"])
    end

    def evaluate_tensor_node(node)
      if !node["shape"].empty? && node["tensor_content"]
        content = node["tensor_content"]
        unpacked = eval(%("#{content}"))

        if node["dtype"] == "DT_FLOAT"
          TensorShape.reshape(unpacked.unpack("f*"), node["shape"])
        elsif node["dtype"] == "DT_INT32"
          TensorShape.reshape(unpacked.unpack("l*"), node["shape"])
        elsif node["dtype"] == "DT_STRING"
          node["string_val"]
        else
          raise "unknown dtype #{node["dtype"]}"
        end
      else

        val = if node["dtype"] == "DT_FLOAT"
          node["float_val"] ? node["float_val"].to_f : []
        elsif node["dtype"] == "DT_INT32"
          node["int_val"] ? node["int_val"].to_i : []
        elsif node["dtype"] == "DT_STRING"
          node["string_val"]
        else
          raise "unknown dtype #{node["dtype"]}"
        end

        if node["shape"] == [1]
          [val]
        else
          val
        end
      end
    end

    def map_type_to_ts(attr_value)
      case attr_value
      when "DT_FLOAT"
        :float32
      when "DT_INT32"
        :int32
      when "DT_INT64"
        :int64
      when "DT_STRING"
        :string
      when "DT_BOOL"
        :boolean
      else
        raise "unknown type #{attr_value}"
      end
    end

    def options_evaluator(node)
      return {} if node["attributes"].nil?

      node["attributes"].map { |attribute|
        attr_type, attr_value = attribute["value"].flat_map { |k, v| [k, v] }

        if attr_type == "tensor"
          attr_value = evaluate_tensor_node(attr_value)
        elsif attr_type == "type"
          attr_value = map_type_to_ts(attr_value)
        elsif attr_type == "b"
          attr_value = attr_value == "true"
        end

        [attribute["key"], attr_value]
      }.to_h
    end

    protected

    def evaluate_lines(lines = [])
      block = []
      node = {}
      node_attr = {}
      state = :top

      lines.each do |str|
        case state
        when :top
          node["type"] = parse_node_name(str)
          state = :node_context
          next
        when :node_context
          if str == "attr {"
            state = :attr_context
            node_attr = {}
            node["attributes"] ||= []
            node["attributes"] << node_attr
            next
          elsif str == "}"
            state = :top
            block << node
            node = {}
            next
          else
            key, value = str.split(":", 2)
            if key == "input"
              node["input"] ||= []
              node["input"] << process_value(value.strip)
            else
              node[key] = process_value(value.strip)
            end
          end
        when :attr_context
          if str == "value {"
            state = :value_context
            node_attr["value"] = {}
            next
          elsif str == "}"
            state = :node_context
            next
          else
            key, value = str.split(":", 2)
            node_attr[key] = process_value(value.strip)
          end
        when :value_context
          if str == "list {"
            state = :list_context
            node_attr["value"] = []
            next
          elsif str == "shape {"
            state = :shape_context
            node_attr["value"]["shape"] = []
            next
          elsif str == "tensor {"
            state = :tensor_context
            node_attr["value"]["tensor"] = {}
            next
          elsif str == "}"
            state = :attr_context
            next
          else
            key, value = str.split(":", 2)
            if key == "dtype"
              node_attr["value"]["dtype"] = value.strip
            elsif key == "type"
              node_attr["value"]["type"] = value.strip
            else
              node_attr["value"][key] = process_value(value.strip)
            end
          end
        when :list_context
          if str == "}"
            state = :value_context
            next
          else
            key, value = str.split(":", 2)
            node_attr["value"] << {key => value}
          end
        when :tensor_context
          if str == "tensor_shape {"
            state = :tensor_shape_context
            node_attr["value"]["tensor"]["shape"] = []
            next
          elsif str == "}"
            state = :value_context
            next
          else
            key, value = str.split(":", 2)
            if node_attr["value"]["tensor"][key] && !node_attr["value"]["tensor"][key].is_a?(Array)
              node_attr["value"]["tensor"][key] = [node_attr["value"]["tensor"][key]]
              node_attr["value"]["tensor"][key] << process_value(value.strip)
            elsif node_attr["value"]["tensor"][key]
              node_attr["value"]["tensor"][key] << process_value(value.strip)
            else
              node_attr["value"]["tensor"][key] = process_value(value.strip)
            end
          end
        when :tensor_shape_context
          if str == "dim {"
            state = :tensor_shape_dim_context
            next
          elsif str == "}"
            state = :tensor_context
            next
          end
        when :shape_context
          if str == "}"
            state = :value_context
            next
          elsif str == "dim {"
            state = :shape_dim_context
            next
          end
        when :shape_dim_context
          if str == "}"
            state = :shape_context
            next
          else
            _key, value = str.split(":", 2)
            node_attr["value"]["shape"] << value.strip.to_i
          end
        when :tensor_shape_dim_context
          if str == "}"
            state = :tensor_shape_context
            next
          else
            _key, value = str.split(":", 2)
            node_attr["value"]["tensor"]["shape"] << value.strip.to_i
          end
        end
      end

      block
    end

    def parse_node_name(str)
      str.split(" ")[0]
    end

    def process_value(value)
      if value.start_with?('"')
        unescape(value.gsub!(/\A"|"\Z/, ""))
      else
        unescape(value)
      end
    end

    UNESCAPES = {
      "a" => "\x07", "b" => "\x08", "t" => "\x09",
      "n" => "\x0a", "v" => "\x0b", "f" => "\x0c",
      "r" => "\x0d", "e" => "\x1b", "\\\\" => "\x5c",
      "\"" => "\x22", "'" => "\x27",
    }.freeze

    def unescape(str)
      # Escape all the things
      str.gsub(/\\(?:([#{UNESCAPES.keys.join}])|u([\da-fA-F]{4}))|\\0?x([\da-fA-F]{2})/) do
        if $1
          $1 == '\\' ? '\\' : UNESCAPES[$1]
        elsif $2 # escape \u0000 unicode
          [$2.to_s.hex].pack("U*")
        elsif $3 # escape \0xff or \xff
          [$3].pack("H2")
        end
      end
    end
  end
end


module TensorStream
  # Class that defines all available ops supported by TensorStream
  module Ops
    if File.exists?(File.join(__dir__, 'generated_stub', 'ops.rb'))
      require 'tensor_stream/generated_stub/ops'
      include TensorStream::OpStub
    end

    class OutputHolder
      def initialize(op)
        @op = op
      end
    end
    FLOATING_POINT_TYPES = %i[float32 float64 float float16].freeze
    INTEGER_TYPES = %i[uint8 int32 int int16 uint16 int64 uint32 uint64].freeze
    NUMERIC_TYPES = FLOATING_POINT_TYPES + INTEGER_TYPES

    ##
    # Assert the condition x == y holds element-wise.
    #
    # Argmuments
    #
    # +x+ Numeric Tensor.
    # +y+ Numeric Tensor, same dtype as and broadcastable to x.
    #
    # Returns
    # Op that raises InvalidArgumentError if x == y is false
    def assert_equal(x, y, data: nil, summarize: nil, message: nil, name: nil)
      _op(:assert_equal, x, y, data: data, summarize: summarize, message: message, name: name)
    end

    ##
    # Constructs symbolic derivatives of ys of input w.r.t. x in wrt_xs.
    #
    # ys and xs are each a Tensor or a list of tensors. grad_ys is a list of Tensor, holding the gradients received by the ys. The list must be the same length as ys.
    #
    # Arguments:
    # +tensor_ys+ : A Tensor or list of tensors to be differentiated.
    # +wrt_xs+ : A Tensor or list of tensors to be used for differentiation.
    # +stop_gradients+ :  Optional. A Tensor or list of tensors not to differentiate through
    def gradients(tensor_ys, wrt_xs, name: "gradients", stop_gradients: nil)
      tensor_ys = tensor_ys.op
      gs = wrt_xs.map(&:op).collect { |x|
        stops = stop_gradients ? stop_gradients.map(&:name).join("_") : ""
        gradient_program_name = "grad_#{tensor_ys.name}_#{x.name}_#{stops}".to_sym
        tensor_graph = tensor_ys.graph

        tensor_program = if tensor_graph.node_added?(gradient_program_name)
          tensor_graph.get_node(gradient_program_name)
        else
          tensor_graph.name_scope("gradient_wrt_#{x.name}") do
            derivative_ops = TensorStream::MathGradients.derivative(tensor_ys, x, graph: tensor_graph,
                                                                                stop_gradients: stop_gradients)
            tensor_graph.add_node!(gradient_program_name, derivative_ops)
          end
        end
        tensor_program
      }

      gs
    end

    ##
    # Outputs random values from a normal distribution.
    def random_normal(shape, dtype: :float32, mean: 0.0, stddev: 1.0, seed: nil, name: nil)
      options = {dtype: dtype, mean: mean, stddev: stddev, seed: seed, name: name}
      _op(:random_standard_normal, shape, options)
    end

    ##
    # Outputs random values from a truncated normal distribution.
    def truncated_normal(shape, dtype: :float32, mean: 0.0, stddev: 1.0, seed: nil, name: nil)
      options = {dtype: dtype, mean: mean, stddev: stddev, seed: seed, name: name}
      _op(:truncated_normal, shape, options)
    end

    ##
    # Stops gradient computation.
    #
    # When executed in a graph, this op outputs its input tensor as-is.
    def stop_gradient(tensor, options = {})
      _op(:stop_gradient, tensor, options)
    end

    ##
    # Construct an identity matrix
    def eye(num_rows, num_columns: nil, dtype: :float32, name: nil)
      _op(:eye, num_rows, num_columns || num_rows, data_type: dtype, name: name)
    end

    def expand_dims(input, axis = nil, name: nil)
      _op(:expand_dims, input, axis, name: name)
    end

    def shape_n(inputs, name: nil, out_type: :int32)
      shapes_known = true
      inputs.each do |input|
        unless input.shape.known?
          shapes_known = false
          break
        end
      end

      if shapes_known
        inputs.collect { |input| cons(input.shape.shape, dtype: out_type).op }
      else
        res = _op(:shape_n, *inputs, out_type: out_type, name: name)
        Array.new(inputs.size) do |index|
          res[index]
        end
      end
    end

    ##
    # initializer that generates tensors initialized to 0.
    #
    def zeros_initializer(dtype: :float32)
      TensorStream::Initializer.new(-> { _op(:zeros, data_type: dtype) })
    end

    ##
    # initializer that generates tensors initialized to 1.
    #
    def ones_initializer(dtype: :float32)
      TensorStream::Initializer.new(-> { _op(:ones, data_type: dtype) })
    end

    def constant_initializer(value, dtype: nil, verify_shape: false)
      TensorStream::Initializer.new(-> { _op(:fill, nil, convert_to_tensor(value, dtype: dtype)) })
    end

    ##
    # The Glorot uniform initializer, also called Xavier uniform initializer.
    #
    # It draws samples from a uniform distribution within [-limit, limit]
    # where limit is sqrt(6 / (fan_in + fan_out)) where fan_in is the number
    # of input units in the weight tensor and fan_out is the number of output units in the weight tensor.
    def glorot_uniform_initializer(seed: nil, dtype: nil)
      TensorStream::Initializer.new(-> { _op(:glorot_uniform, seed: seed, data_type: dtype) })
    end

    ##
    # Initializer that generates tensors with a uniform distribution.
    def random_uniform_initializer(minval: 0, maxval: 1, seed: nil, dtype: nil)
      TensorStream::Initializer.new(-> { _op(:random_uniform, minval: 0, maxval: 1, seed: seed, data_type: dtype) })
    end

    ##
    # Extracts a slice from a tensor.
    #
    # This operation extracts a slice of size size from a tensor input starting at the location specified by begin.
    # The slice size is represented as a tensor shape, where size[i] is the number of elements of the 'i'th dimension of input that you want to slice. The starting location (begin) for the slice is
    # represented as an offset in each dimension of input. In other words, begin[i] is the offset into the 'i'th dimension of input that you want to slice from.
    def slice(input, start, size, name: nil)
      _op(:slice, input, start, size: size, name: name)
    end

    ##
    # Creates a tensor with all elements set to 1.
    def ones(shape, dtype: :float32, name: nil)
      _op(:ones, shape, data_type: dtype, name: name)
    end

    ##
    # Returns the truth value of x AND y element-wise.
    def logical_and(input_a, input_b, name: nil)
      check_data_types(input_a, input_b)
      _op(:logical_and, input_a, input_b, name: name)
    end

    ##
    # Computes the mean of elements across dimensions of a tensor.
    def reduce_mean(input_tensor, axis = nil, keepdims: false, name: nil)
      reduce(:mean, input_tensor, axis, keepdims: keepdims, name: name)
    end

    def reduce(op, input, axis = nil, keepdims: false, name: nil)
      input = TensorStream.convert_to_tensor(input)
      return input if input.shape.scalar?

      axis = cast_axis(input, axis)

      _op(op, input, axis, keepdims: keepdims, name: name)
    end

    ##
    # Concatenates tensors along one dimension.
    def concat(values, axis, name: "concat")
      if values.is_a?(Array)
        _op(:concat, axis, *values, name: name)
      else
        _op(:concat, axis, values, name: name)
      end
    end

    ##
    # Partitions data into num_partitions tensors using indices from partitions
    def dynamic_partition(data, partitions, num_partitions, name: nil)
      result = _op(:dynamic_partition, data, partitions, num_partitions: num_partitions, name: nil)
      num_partitions.times.map do |index|
        result[index]
      end
    end

    def split(value, num_or_size_splits, axis: 0, num: nil, name: "split")
      value = convert_to_tensor(value)
      num_or_size_splits = convert_to_tensor(num_or_size_splits)
      axis = convert_to_tensor(axis)

      raise TensorStream::ValueError, "num_or_size_splits must be integer dtype" unless INTEGER_TYPES.include?(num_or_size_splits.data_type)

      res = _op(:split, value, num_or_size_splits, axis, name: name)

      pieces = if value.shape.known? && num_or_size_splits.is_const && num_or_size_splits.value && axis.is_const
        if num_or_size_splits.shape.scalar?
          raise TensorStream::ValueError, "num_or_size_splits must divide dimension #{value.shape.shape[axis.value]} evenly" unless (value.shape.shape[axis.value] % num_or_size_splits.value).zero?

          div = num_or_size_splits.value
          n = value.shape.shape[axis.value] / div

          Array.new(div) do
            new_shape = value.shape.shape.dup
            new_shape[axis.value] = n
            new_shape
          end
        elsif num_or_size_splits.shape.ndims == 1
          raise TensorStream::ValueError, "Sum of splits do not match total dimen in axis #{value.shape.shape[axis.value]} != #{num_or_size_splits.value.reduce(:+)}" if value.shape.shape[axis.value] != num_or_size_splits.value.reduce(:+)

          num_or_size_splits.value.collect do |v|
            new_shape = value.shape.shape.dup
            new_shape[axis.value] = v
            new_shape
          end
        else
          raise TensorStream::ValueError, "Scalar or 1D Tensor expected for num_or_size_splits"
        end
      else
        raise TensorStream::ValueError, "Cannot automatically determine num, please specify num: in options" if num.nil?

        Array.new(num) { nil }
      end

      pieces.collect.with_index do |shape, i|
        op = index(res, i, name: "split/index:#{i}")
        op.shape = TensorShape.new(shape) if shape

        op
      end
    end

    ##
    # select an index in an array or a set of tensor outputs
    def index(tensor, sel, name: nil)
      _op(:index, tensor, sel, name: name)
    end

    ##
    # Reshapes a tensor.
    #
    # Given tensor, this operation returns a tensor that has the same values as tensor with shape shape.
    def reshape(tensor, shape, name: nil)
      _op(:reshape, tensor, shape, name: name)
    end

    ##
    # Computes square of x element-wise.
    def square(tensor, name: nil)
      _op(:square, tensor, name: name)
    end

    ##
    # Computes the reciprocal of x element-wise.
    def reciprocal(tensor, name: nil)
      _op(:reciprocal, tensor, name: name)
    end

    ##
    # Return true_fn() if the predicate pred is true else false_fn().
    def cond(pred, true_fn, false_fn, name: nil)
      _op(:case, [pred], false_fn, true_fn, name: name)
    end

    ##
    # Return the elements, either from x or y, depending on the condition.
    def where(condition, true_t = nil, false_t = nil, name: nil)
      _op(:where, condition, true_t, false_t, name: name)
    end

    ##
    # Adds all input tensors element-wise.
    #
    # Elements must all be the same shape and type
    def add_n(inputs, name: nil)
      _op(:add_n, *inputs, name: name)
    end

    ##
    # Computes asin of input element-wise
    def asin(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:asin, input, name: name)
    end

    ##
    # Computes acos of input element-wise
    def acos(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:acos, input, name: name)
    end

    ##
    # Computes atan of input element-wise
    def atan(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:atan, input, name: name)
    end

    ##
    # Returns element-wise integer divistion.
    def floor_div(input_a, input_b, name: nil)
      check_data_types(input_a, input_b)
      _op(:floor_div, input_a, input_b, name: name)
    end

    ##
    # Casts a tensor to a new type, if needed
    def cast(input, dtype, name: nil)
      input = convert_to_tensor(input)
      return input if input.data_type == dtype

      _op(:cast, input, data_type: dtype, name: name)
    end

    ##
    # Returns the max of x and y (i.e. x > y ? x : y) element-wise.
    def maximum(input_a, input_b, name: nil)
      max(input_a, input_b, name: name)
    end

    ##
    # Returns the min of x and y (i.e. x < y ? x : y) element-wise.
    def minimum(input_a, input_b, name: nil)
      min(input_a, input_b, name: name)
    end

    ##
    # Prints a list of tensors.
    #
    # This is an identity op (behaves like tf.identity) with the side effect of printing data when evaluating.
    def print(input, data, message: nil, name: nil)
      _op(:print, input, data, message: message, name: name)
    end

    ##
    # Computes numerical negative value element-wise.
    def negative(input, name: nil)
      negate(input, name: name)
    end

    ##
    # Returns the truth value of (x != y) element-wise.
    # This ops supports broadcasting
    def not_equal(input_a, input_b, name: nil)
      check_data_types(input_a, input_b)
      _op(:not_equal, input_a, input_b, name: name)
    end

    ##
    # reates a tensor with all elements set to zero.
    # Given a single tensor (tensor), this operation returns a tensor
    # of the same type and shape as tensor with all elements set to zero.
    # Optionally, you can use dtype to specify a new type for the returned tensor.
    def zeros_like(tensor, dtype: nil, name: nil)
      _op(:zeros_like, tensor, data_type: dtype, name: name)
    end

    ##
    # Creates a tensor with all elements set to 1.
    # Given a single tensor (tensor), this operation returns a
    # tensor of the same type and shape as tensor with all elements set to 1.
    # Optionally, you can specify a new type (dtype) for the returned tensor.
    def ones_like(tensor, dtype: nil, name: nil)
      _op(:ones_like, tensor, data_type: dtype, name: name)
    end

    ##
    # Return a tensor with the same shape and contents as input.
    def identity(input, name: nil)
      _op(:identity, input, name: name)
    end

    ##
    # Returns x * y element-wise.
    # This operation supports broadcasting
    def multiply(input_a, input_b, name: nil)
      check_data_types(input_a, input_b)
      _op(:mul, input_a, input_b, name: name)
    end

    ##
    # Computes the absolute value of a tensor.
    def abs(input, name: nil)
      _op(:abs, input, name: name)
    end

    ##
    # Computes sec of input element-wise.
    def sec(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:sec, input, name: name)
    end

    ##
    # Computes sqrt of input element-wise.
    def sqrt(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:sqrt, input, name: name)
    end

    ##
    # Computes natural logarithm of x element-wise.
    def log(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:log, input, name: name)
    end

    ##
    # Computes natural logarithm of (1 + x) element-wise.
    def log1p(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:log1p, input, name: name)
    end

    ##
    # Computes exponential of x element-wise.
    def exp(input, name: nil)
      check_allowed_types(input, FLOATING_POINT_TYPES)
      _op(:exp, input, name: name)
    end

    ##
    # Transposes a. Permutes the dimensions according to perm.
    def transpose(tensor, perm = nil, name: "transpose")
      _op(:transpose, tensor, perm, name: name)
    end

    ##
    # Pads a tensor.
    # This operation pads a tensor according to the paddings you specify.
    def pad(tensor, paddings, mode: "CONSTANT", name: nil)
      _op(:pad, tensor, paddings, mode: mode, name: name)
    end

    ##
    # Checks a tensor for NaN and Inf values.
    # When run, reports an InvalidArgument error if tensor has any values that are not a number (NaN) or infinity (Inf). Otherwise, passes tensor as-is.
    def check_numerics(tensor, message, name: nil)
      _op(:check_numerics, tensor, message: message, name: name)
    end

    def squared_difference(input_a, input_b, name: nil)
      _op(:squared_difference, input_a, input_b, name: name)
    end

    def broadcast_gradient_args(shape_a, shape_b, name: nil)
      op_result = _op(:broadcast_gradient_args, shape_a, shape_b, name: name)
      [op_result[0], op_result[1]]
    end

    ##
    # Gather slices from params and axis according to indices.
    #
    def gather(params, indices, validate_indices: nil,
      name: nil,
      axis: 0)
      _op(:gather, params, indices, validate_indices: validate_indices, name: name, axis: axis)
    end

    ##
    # Stacks a list of rank-R tensors into one rank-(R+1) tensor.
    #
    def stack(values, axis: 0, name: "stack")
      _op(:stack, *values, axis: axis, name: name)
    end

    ##
    # Unpacks the given dimension of a rank-R tensor into rank-(R-1) tensors.
    #
    def unstack(value, num: nil, axis: 0, name: "unstack")
      res = _op(:unstack, value, num: num, axis: axis, name: name)

      num_vars = if value.shape.known?
        new_shape = value.shape.shape.dup
        rank = new_shape.size - 1
        axis = rank + axis if axis < 0
        rotated_shape = Array.new(axis + 1) { new_shape.shift }
        new_shape = rotated_shape.rotate!(-1) + new_shape
        new_shape[0]
      else
        raise TensorStream::ValueError, "num is unspecified and cannot be inferred." if num.nil?

        num
      end

      return res[0] if num_vars == 1

      Array.new(num_vars) do |i|
        index(res, i, name: "unstack/index:#{i}")
      end
    end

    ##
    # Same as stack
    def pack(values, axis: 0, name: "pack")
      _op(:stack, *values, axis: axis, name: name)
    end

    ##
    # Same as unstack
    #
    def unpack(value, num: nil, axis: 0, name: "unpack")
      unstack(value, num: num, axis: axis, name: name)
    end

    ##
    # Removes dimensions of size 1 from the shape of a tensor.
    #
    # Given a tensor input, this operation returns a tensor of the same type with all dimensions of size 1 removed.
    # If you don't want to remove all size 1 dimensions, you can remove specific size 1 dimensions by specifying axis.
    def squeeze(value, axis: [], name: nil)
      _op(:squeeze, value, axis: axis, name: nil)
    end

    def clip_by_norm(tensor, clip_norm, axes: nil, name: nil)
    end

    ##
    # Computes the difference between two lists of numbers or strings.
    # Given a list x and a list y, this operation returns a list out that represents all values
    # that are in x but not in y. The returned list out is sorted in the same order that the numbers appear
    # in x (duplicates are preserved). This operation also returns a list idx that represents the position of
    # each out element in x. In other words:
    #
    def setdiff1d(x, y, index_dtype: :int32, name: nil)
      result = _op(:setdiff1d, x, y, index_dtype: index_dtype, name: name)
      [result[0], result[1]]
    end

    ##
    # Create a case operation.
    #
    # The pred_fn_pairs parameter is a dict or list of pairs of size N.
    # Each pair contains a boolean scalar tensor and a proc that creates the tensors to be returned if the boolean evaluates to true.
    # default is a proc generating a list of tensors. All the proc in pred_fn_pairs as well as default (if provided) should return the
    # same number and types of tensors.
    #
    def case(args = {})
      args = args.dup
      default = args.delete(:default)
      exclusive = args.delete(:exclusive)
      strict = args.delete(:strict)
      name = args.delete(:name)

      predicates = []
      functions = []

      args.each do |k, v|
        raise "Invalid argment or option #{k}" unless k.is_a?(Tensor)

        predicates << k
        functions << (v.is_a?(Proc) ? v.call : v)
      end

      _op(:case, predicates, default, *functions, exclusive: exclusive, strict: strict, name: name)
    end

    def cumprod(x, axis: 0, exclusive: false, reverse: false, name: nil)
      _op(:cumprod, x, axis: axis, exclusive: exclusive, reverse: reverse, name: name)
    end

    def invert_permutation(x, name: nil)
      _op(:invert_permutation, x, name: name)
    end

    def cast_axis(input, axis)
      if !axis.nil?
        axis
      elsif input.shape.known?
        (0...input.shape.ndims).to_a
      else
        range(0, rank(input))
      end
    end
  end
end


module TensorStream
  # module that contains helper functions useful for ops
  module OpHelper
    def _op(code, *args)
      default_graph = Graph.get_default_graph

      op = default_graph.add_op!(code.to_sym, *args)
      if !default_graph.get_dependency_scope.nil?
        i_op(:identity, op, default_graph.get_dependency_scope, name: [op.name, "tuple", "control_dependency"].join("/"))
      else
        op
      end
    end

    # same as op but with a marker that it was internal generated
    def i_op(code, *args)
      options = if args.last.is_a?(Hash)
        args.pop
      else
        {}
      end

      args << options.merge(internal: true)
      Graph.get_default_graph.add_op!(code.to_sym, *args)
    end

    def i_var(data_type, rank, shape, variable_scope, options = {})
      new_var = Variable.new(data_type)
      new_var.prepare(rank, shape, variable_scope, options)
      new_var.op = new_var.graph.add_variable!(new_var, options.merge(shape: @shape, data_type: data_type))

      new_var
    end

    def cons(value, options = {})
      TensorStream.constant(value, options)
    end

    def i_cons(value, options = {})
      TensorStream.constant(value, options.merge(internal: true))
    end

    def shape_eval(input, output_type = :int32)
      return [] unless input.is_a?(Array)
      arr = []
      arr_ptr = input

      Kernel.loop do
        arr << (TensorStream::Ops::FLOATING_POINT_TYPES.include?(output_type) ? arr_ptr.size.to_f : arr_ptr.size)
        arr_ptr = arr_ptr[0]

        break unless arr_ptr.is_a?(Array)
      end

      arr
    end

    def fp_type?(type)
      TensorStream::Ops::FLOATING_POINT_TYPES.include?(type)
    end

    def int_type?(type)
      TensorStream::Ops::INTEGER_TYPES.include?(type)
    end

    def format_source(trace)
      grad_source = trace.detect { |c| c.to_s.include?(File.join("lib", "tensor_stream", "math_gradients")) }
      source = trace.reject { |c| c.to_s.include?(File.join("lib", "tensor_stream")) }.first
      [grad_source, trace].compact.join("\n")
    end

    def shapes_fully_specified_and_equal(x, y)
      return false if !shape_full_specified(x) || !shape_full_specified(y)
      return false if x.shape.shape != y.shape.shape

      true
    end

    def shape_full_specified(tensor)
      return false if tensor.shape.nil?
      return false if tensor.shape.shape.nil?

      tensor.shape.shape.each { |s| return false if s.nil? || (s < 0) }
      true
    end

    def reduced_shape(input_shape, axes)
      input_shape = TensorStream.convert_to_tensor(input_shape)
      axes = TensorStream.convert_to_tensor(axes)
      input_rank = i_op(:size, input_shape)
      axes = TensorStream.range(0, input_rank) if axes.nil?
      axes = (axes + input_rank) % input_rank
      axes_shape = i_op(:shape, axes)

      TensorStream.dynamic_stitch([TensorStream.range(0, input_rank), axes],
        [input_shape, i_op(:fill, axes_shape, 1)])
    end
  end
end


module TensorStream
  # varoius utility functions for array processing
  module ArrayOpsHelper
    def split_tensor(input, begin_index, end_index, axis = 0)
      if axis.zero?
        input[begin_index...end_index]
      else
        input.collect do |item|
          split_tensor(item, begin_index, end_index, axis - 1)
        end
      end
    end

    def slice_tensor(input, start, size)
      return input if size.empty?
      start_index = start.shift
      current_size = size.shift
      dimen_size = if current_size == -1
        input.size - 1
      else
        start_index + current_size - 1
      end

      input[start_index..dimen_size].collect do |item|
        if item.is_a?(Array)
          slice_tensor(item, start.dup, size.dup)
        else
          item
        end
      end
    end

    def array_set!(input, value)
        input.each_with_index do |element, index|
          if element.is_a?(Array)
            array_set(element, value)
          else
            input[index] = value[index]
          end
        end
    end

    def truncate(input, target_shape)
      rank = get_rank(input)
      return input if rank.zero?

      start = Array.new(rank) { 0 }
      slice_tensor(input, start, target_shape)
    end

    def _reduced_shape(input_shape, axes)
      return [] if axes.nil? # reduce to scalar
      axes = [axes] unless axes.is_a?(Array)
      return input_shape if axes.empty?

      axes.each do |dimen|
        input_shape[dimen] = 1
      end
      input_shape
    end

    def broadcast(input_a, input_b)
      sa = shape_eval(input_a)
      sb = shape_eval(input_b)

      return [input_a, input_b] if sa == sb

      # descalar
      if sa.empty?
        input_a = [input_a]
        sa = [1]
      end

      if sb.empty?
        input_b = [input_b]
        sb = [1]
      end

      target_shape = shape_diff(sa, sb)

      if target_shape
        input_b = broadcast_dimensions(input_b, target_shape)
      else
        target_shape = shape_diff(sb, sa)
        raise "Incompatible shapes for op #{shape_eval(input_a)} vs #{shape_eval(input_b)}" if target_shape.nil?

        input_a = broadcast_dimensions(input_a, target_shape)
      end

      [input_a, input_b]
    end

    # explicit broadcasting helper
    def broadcast_dimensions(input, dims = [])
      return input if dims.empty?

      d = dims.shift

      if input.is_a?(Array) && (get_rank(input) - 1) == dims.size
        row_to_dup = input.collect { |item|
          broadcast_dimensions(item, dims.dup)
        }

        row_to_dup + Array.new(d) { row_to_dup }.flatten(1)
      elsif input.is_a?(Array)
        Array.new(d) { broadcast_dimensions(input, dims.dup) }
      else
        Array.new(d + 1) { input }
      end
    end

    # handle 2 tensor math operations
    def vector_op(vector, vector2, switch = false, safe = true, &block)
      if get_rank(vector) < get_rank(vector2) # upgrade rank of A
        duplicated = Array.new(vector2.size) {
          vector
        }
        return vector_op(duplicated, vector2, switch, &block)
      end

      return yield(vector, vector2) unless vector.is_a?(Array)

      vector.each_with_index.collect { |input, index|
        next vector_op(input, vector2, switch, &block) if input.is_a?(Array) && get_rank(vector) > get_rank(vector2)

        if safe && vector2.is_a?(Array)
          next nil if vector2.size != 1 && index >= vector2.size
        end

        z = if vector2.is_a?(Array)
          if index < vector2.size
            vector2[index]
          else
            raise "incompatible tensor shapes used during op" if vector2.size != 1
            vector2[0]
          end
        else
          vector2
        end

        if input.is_a?(Array)
          vector_op(input, z, switch, &block)
        else
          switch ? yield(z, input) : yield(input, z)
        end
      }.compact
    end

    def shape_diff(shape_a, shape_b)
      return nil if shape_b.size > shape_a.size

      reversed_a = shape_a.reverse
      reversed_b = shape_b.reverse

      reversed_a.each_with_index.collect { |s, index|
        next s if index >= reversed_b.size
        return nil if reversed_b[index] > s
        s - reversed_b[index]
      }.reverse
    end

    def tile_arr(input, dimen, multiples)
      t = multiples[dimen]
      if dimen == multiples.size - 1
        return nil if t.zero?
        input * t # ruby array dup
      else
        new_arr = input.collect { |sub|
          tile_arr(sub, dimen + 1, multiples)
        }.compact

        return nil if new_arr.empty?

        new_arr * t
      end
    end

    def process_function_op(a, &block)
      # ruby scalar
      if (a.is_a?(Tensor) && a.shape.rank > 0) || a.is_a?(Array)
        vector_op(a, 0, &block)
      else
        yield a, 0
      end
    end

    def get_rank(value, rank = 0)
      return rank unless value.is_a?(Array)
      return rank + 1 if value.empty?

      get_rank(value[0], rank + 1)
    end

    def last_axis(arr)
      return arr if get_rank(arr) <= 2

      arr.inject([]).map do |sub, rows|
        rows + last_axis(sub)
      end
    end

    def softmax(arr)
      return arr if arr.empty?

      if !arr[0].is_a?(Array)
        c = arr.max
        arr = arr.map { |a| Math.exp(a - c) }
        sum = arr.reduce(:+)
        arr.collect do |input|
          input / sum
        end
      else
        arr.collect { |input| softmax(input) }
      end
    end

    def softmax_grad(arr)
      return arr if arr.empty?
      arr.each_with_index.collect do |input, index|
        if input.is_a?(Array)
          softmax_grad(input)
        else
          arr.each_with_index.collect do |input2, index2|
            if index != index2
              -input * input2
            else
              input * (1.0 - input)
            end
          end
        end
      end
    end

    def gather(params, indexes)
      indexes.collect do |index|
        if index.is_a?(Array)
          gather(params, index)
        else
          params[index]
        end
      end
    end

    # general case transposition with flat arrays
    def transpose_with_perm(arr, new_arr, shape, new_shape, perm)
      arr_size = shape.reduce(:*)
      divisors = shape.dup.drop(1).reverse.inject([1]) { |a, s|
        a << s * a.last
      }.reverse

      multipliers = new_shape.dup.drop(1).reverse.inject([1]) { |a, s|
        a << s * a.last
      }.reverse

      arr_size.times do |p|
        ptr = p
        index = []
        divisors.each_with_object(index) do |div, a|
          a << (ptr / div.to_f).floor
          ptr = ptr % div
        end

        # remap based on perm
        remaped = perm.map { |x| index[x] }

        ptr2 = 0
        multipliers.each_with_index do |m, idx|
          ptr2 += remaped[idx] * m
        end

        new_arr[ptr2] = arr[p]
      end

      [new_arr, new_shape]
    end

    def reduce_axis(current_axis, axis, val, keep_dims, &block)
      return val unless val.is_a?(Array)

      r = val.collect { |v|
        reduce_axis(current_axis + 1, axis, v, keep_dims, &block)
      }

      should_reduce_axis = axis.nil? || (axis.is_a?(Array) && axis.include?(current_axis)) || (current_axis == axis)

      if should_reduce_axis
        reduced_val = r[0]
        if r.size > 1
          reduced_val = if block_given?
            yield(r[0..val.size])
          else
            r[0..val.size].reduce(:+)
          end
        elsif r.empty?
          reduced_val = yield(nil)
        end
        keep_dims ? [reduced_val] : reduced_val
      else
        r
      end
    end

    def reduce(val, axis, keep_dims, &block)
      rank = get_rank(val)
      return val if axis&.is_a?(Array) && axis&.empty?

      axis = if axis.nil?
        nil
      elsif axis.is_a?(Array)
        return val if axis.empty?

        axis.map { |a| a < 0 ? rank - a.abs : a }
      else
        axis < 0 ? rank - axis.abs : axis
      end

      reduce_axis(0, axis, val, keep_dims, &block)
    end

    def arr_pad(arr, paddings, data_type = :float32, rank = 0)
      raise "padding #{paddings[rank]} needs to have to elements [before, after]" if paddings[rank].size != 2

      before = paddings[rank][0]
      after = paddings[rank][1]
      pad_value = fp_type?(data_type) ? 0.0 : 0
      if arr[0].is_a?(Array)
        next_dim_elem = arr.collect { |a| arr_pad(a, paddings, data_type, rank + 1) }
        padding = deep_dup_array(next_dim_elem[0], pad_value)
        Array.new(before) { padding } + next_dim_elem + Array.new(after) { padding }
      else
        Array.new(before) { pad_value } + arr + Array.new(after) { pad_value }
      end
    end

    def deep_dup_array(arr, value = nil)
      if arr.is_a?(Array)
        arr.dup.collect do |a|
          deep_dup_array(a, value)
        end
      else
        value.nil? ? arr : value
      end
    end

    def strided_slice(value, slices = [])
      current_slice = slices.dup
      selection = current_slice.shift
      return value if selection.nil?

      b, e, stride = selection

      b = value.size + b if b < 0
      e = value.size + e + 1 if e < 0

      indexes = if stride < 0
                  b.downto(e).select.with_index { |elem, index| (index % stride.abs) == 0 }
                else
                  (b...e).step(stride)
                end

      indexes.map do |index|
        strided_slice(value[index], current_slice)
      end
    end

    def strided_slice_grad(value, grad, x, slices)
      current_slice = slices.dup
      selection = current_slice.shift
      current_shape = x.shift

      if selection.nil?
        array_set!(value, grad)
      end

      b, e, stride = selection

      b = value.size + b if b < 0
      e = value.size + e + 1 if e < 0

      indexes = if stride < 0
                  b.downto(e).select.with_index { |elem, index| (index % stride.abs) == 0 }
                else
                  (b...e).step(stride)
                end

      indexes.each_with_index do |index, grad_index|
        if (value[index].is_a?(Array))
          strided_slice_grad(value[index], grad[grad_index], x.dup, current_slice.dup)
        else
          value[index] = grad[grad_index]
        end
      end
    end
  end
end