/*
 * Copyright (C) 2015 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.squareup.sqlbrite3;

import android.arch.persistence.db.SimpleSQLiteQuery;
import android.arch.persistence.db.SupportSQLiteDatabase;
import android.arch.persistence.db.SupportSQLiteOpenHelper;
import android.arch.persistence.db.SupportSQLiteOpenHelper.Callback;
import android.arch.persistence.db.SupportSQLiteQuery;
import android.arch.persistence.db.SupportSQLiteStatement;
import android.content.ContentValues;
import android.database.Cursor;
import android.database.sqlite.SQLiteTransactionListener;
import android.support.annotation.CheckResult;
import android.support.annotation.IntDef;
import android.support.annotation.NonNull;
import android.support.annotation.Nullable;
import android.support.annotation.WorkerThread;
import com.squareup.sqlbrite3.SqlBrite.Logger;
import com.squareup.sqlbrite3.SqlBrite.Query;
import io.reactivex.Observable;
import io.reactivex.ObservableTransformer;
import io.reactivex.Scheduler;
import io.reactivex.functions.Consumer;
import io.reactivex.functions.Function;
import io.reactivex.functions.Predicate;
import io.reactivex.subjects.PublishSubject;
import io.reactivex.subjects.Subject;
import java.io.Closeable;
import java.lang.annotation.Retention;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import static android.database.sqlite.SQLiteDatabase.CONFLICT_ABORT;
import static android.database.sqlite.SQLiteDatabase.CONFLICT_FAIL;
import static android.database.sqlite.SQLiteDatabase.CONFLICT_IGNORE;
import static android.database.sqlite.SQLiteDatabase.CONFLICT_NONE;
import static android.database.sqlite.SQLiteDatabase.CONFLICT_REPLACE;
import static android.database.sqlite.SQLiteDatabase.CONFLICT_ROLLBACK;
import static com.squareup.sqlbrite3.QueryObservable.QUERY_OBSERVABLE;
import static java.lang.annotation.RetentionPolicy.SOURCE;
import static java.util.Collections.singletonList;

/**
 * A lightweight wrapper around {@link SupportSQLiteOpenHelper} which allows for continuously
 * observing the result of a query. Create using a {@link SqlBrite} instance.
 */
public final class BriteDatabase implements Closeable {
  private final SupportSQLiteOpenHelper helper;
  private final Logger logger;
  private final ObservableTransformer<Query, Query> queryTransformer;

  // Package-private to avoid synthetic accessor method for 'transaction' instance.
  final ThreadLocal<SqliteTransaction> transactions = new ThreadLocal<>();
  private final Subject<Set<String>> triggers = PublishSubject.create();

  private final Transaction transaction = new Transaction() {
    @Override public void markSuccessful() {
      if (logging) log("TXN SUCCESS %s", transactions.get());
      getWritableDatabase().setTransactionSuccessful();
    }

    @Override public boolean yieldIfContendedSafely() {
      return getWritableDatabase().yieldIfContendedSafely();
    }

    @Override public boolean yieldIfContendedSafely(long sleepAmount, TimeUnit sleepUnit) {
      return getWritableDatabase().yieldIfContendedSafely(sleepUnit.toMillis(sleepAmount));
    }

    @Override public void end() {
      SqliteTransaction transaction = transactions.get();
      if (transaction == null) {
        throw new IllegalStateException("Not in transaction.");
      }
      SqliteTransaction newTransaction = transaction.parent;
      transactions.set(newTransaction);
      if (logging) log("TXN END %s", transaction);
      getWritableDatabase().endTransaction();
      // Send the triggers after ending the transaction in the DB.
      if (transaction.commit) {
        sendTableTrigger(transaction);
      }
    }

    @Override public void close() {
      end();
    }
  };
  private final Consumer<Object> ensureNotInTransaction = new Consumer<Object>() {
    @Override public void accept(Object ignored) throws Exception {
      if (transactions.get() != null) {
        throw new IllegalStateException("Cannot subscribe to observable query in a transaction.");
      }
    }
  };

  private final Scheduler scheduler;

  // Package-private to avoid synthetic accessor method for 'transaction' instance.
  volatile boolean logging;

  BriteDatabase(SupportSQLiteOpenHelper helper, Logger logger, Scheduler scheduler,
      ObservableTransformer<Query, Query> queryTransformer) {
    this.helper = helper;
    this.logger = logger;
    this.scheduler = scheduler;
    this.queryTransformer = queryTransformer;
  }

  /**
   * Control whether debug logging is enabled.
   */
  public void setLoggingEnabled(boolean enabled) {
    logging = enabled;
  }

  /**
   * Create and/or open a database.  This will be the same object returned by
   * {@link SupportSQLiteOpenHelper#getWritableDatabase} unless some problem, such as a full disk,
   * requires the database to be opened read-only.  In that case, a read-only
   * database object will be returned.  If the problem is fixed, a future call
   * to {@link SupportSQLiteOpenHelper#getWritableDatabase} may succeed, in which case the read-only
   * database object will be closed and the read/write object will be returned
   * in the future.
   *
   *
Like {@link SupportSQLiteOpenHelper#getWritableDatabase}, this method may
   * take a long time to return, so you should not call it from the
   * application main thread, including from
   * {@link android.content.ContentProvider#onCreate ContentProvider.onCreate()}.
   *
   * @throws android.database.sqlite.SQLiteException if the database cannot be opened
   * @return a database object valid until {@link SupportSQLiteOpenHelper#getWritableDatabase}
   *     or {@link #close} is called.
   */
  @NonNull @CheckResult @WorkerThread
  public SupportSQLiteDatabase getReadableDatabase() {
    return helper.getReadableDatabase();
  }

  /**
   * Create and/or open a database that will be used for reading and writing.
   * The first time this is called, the database will be opened and
   * {@link Callback#onCreate}, {@link Callback#onUpgrade} and/or {@link Callback#onOpen} will be
   * called.
   *
   *
Once opened successfully, the database is cached, so you can
   * call this method every time you need to write to the database.
   * (Make sure to call {@link #close} when you no longer need the database.)
   * Errors such as bad permissions or a full disk may cause this method
   * to fail, but future attempts may succeed if the problem is fixed.

   *
   *
Database upgrade may take a long time, you
   * should not call this method from the application main thread, including
   * from {@link android.content.ContentProvider#onCreate ContentProvider.onCreate()}.
   *
   * @throws android.database.sqlite.SQLiteException if the database cannot be opened for writing
   * @return a read/write database object valid until {@link #close} is called
   */
  @NonNull @CheckResult @WorkerThread
  public SupportSQLiteDatabase getWritableDatabase() {
    return helper.getWritableDatabase();
  }

  void sendTableTrigger(Set<String> tables) {
    SqliteTransaction transaction = transactions.get();
    if (transaction != null) {
      transaction.addAll(tables);
    } else {
      if (logging) log("TRIGGER %s", tables);
      triggers.onNext(tables);
    }
  }

  /**
   * Begin a transaction for this thread.
   *

   * Transactions may nest. If the transaction is not in progress, then a database connection is
   * obtained and a new transaction is started. Otherwise, a nested transaction is started.
   *

   * Each call to {@code newTransaction} must be matched exactly by a call to
   * {@link Transaction#end()}. To mark a transaction as successful, call
   * {@link Transaction#markSuccessful()} before calling {@link Transaction#end()}. If the
   * transaction is not successful, or if any of its nested transactions were not successful, then
   * the entire transaction will be rolled back when the outermost transaction is ended.
   *

   * Transactions queue up all query notifications until they have been applied.
   *

   * Here is the standard idiom for transactions:
   *
   *
{@code
   * try (Transaction transaction = db.newTransaction()) {
   *   ...
   *   transaction.markSuccessful();
   * }
   * }

   *
   * Manually call {@link Transaction#end()} when try-with-resources is not available:
   *
{@code
   * Transaction transaction = db.newTransaction();
   * try {
   *   ...
   *   transaction.markSuccessful();
   * } finally {
   *   transaction.end();
   * }
   * }

   *
   *
   * @see SupportSQLiteDatabase#beginTransaction()
   */
  @CheckResult @NonNull
  public Transaction newTransaction() {
    SqliteTransaction transaction = new SqliteTransaction(transactions.get());
    transactions.set(transaction);
    if (logging) log("TXN BEGIN %s", transaction);
    getWritableDatabase().beginTransactionWithListener(transaction);

    return this.transaction;
  }

  /**
   * Begins a transaction in IMMEDIATE mode for this thread.
   *

   * Transactions may nest. If the transaction is not in progress, then a database connection is
   * obtained and a new transaction is started. Otherwise, a nested transaction is started.
   *

   * Each call to {@code newNonExclusiveTransaction} must be matched exactly by a call to
   * {@link Transaction#end()}. To mark a transaction as successful, call
   * {@link Transaction#markSuccessful()} before calling {@link Transaction#end()}. If the
   * transaction is not successful, or if any of its nested transactions were not successful, then
   * the entire transaction will be rolled back when the outermost transaction is ended.
   *

   * Transactions queue up all query notifications until they have been applied.
   *

   * Here is the standard idiom for transactions:
   *
   *
{@code
   * try (Transaction transaction = db.newNonExclusiveTransaction()) {
   *   ...
   *   transaction.markSuccessful();
   * }
   * }

   *
   * Manually call {@link Transaction#end()} when try-with-resources is not available:
   *
{@code
   * Transaction transaction = db.newNonExclusiveTransaction();
   * try {
   *   ...
   *   transaction.markSuccessful();
   * } finally {
   *   transaction.end();
   * }
   * }

   *
   *
   * @see SupportSQLiteDatabase#beginTransactionNonExclusive()
   */
  @CheckResult @NonNull
  public Transaction newNonExclusiveTransaction() {
    SqliteTransaction transaction = new SqliteTransaction(transactions.get());
    transactions.set(transaction);
    if (logging) log("TXN BEGIN %s", transaction);
    getWritableDatabase().beginTransactionWithListenerNonExclusive(transaction);

    return this.transaction;
  }

  /**
   * Close the underlying {@link SupportSQLiteOpenHelper} and remove cached readable and writeable
   * databases. This does not prevent existing observables from retaining existing references as
   * well as attempting to create new ones for new subscriptions.
   */
  @Override public void close() {
    helper.close();
  }

  /**
   * Create an observable which will notify subscribers with a {@linkplain Query query} for
   * execution. Subscribers are responsible for always closing {@link Cursor} instance
   * returned from the {@link Query}.
   *

   * Subscribers will receive an immediate notification for initial data as well as subsequent
   * notifications for when the supplied {@code table}'s data changes through the {@code insert},
   * {@code update}, and {@code delete} methods of this class. Unsubscribe when you no longer want
   * updates to a query.
   *

   * Since database triggers are inherently asynchronous, items emitted from the returned
   * observable use the {@link Scheduler} supplied to {@link SqlBrite#wrapDatabaseHelper}. For
   * consistency, the immediate notification sent on subscribe also uses this scheduler. As such,
   * calling {@link Observable#subscribeOn subscribeOn} on the returned observable has no effect.
   *

   * Note: To skip the immediate notification and only receive subsequent notifications when data
   * has changed call {@code skip(1)} on the returned observable.
   *

   * Warning: this method does not perform the query! Only by subscribing to the returned
   * {@link Observable} will the operation occur.
   *
   * @see SupportSQLiteDatabase#query(String, Object[])
   */
  @CheckResult @NonNull
  public QueryObservable createQuery(@NonNull final String table, @NonNull String sql,
      @NonNull Object... args) {
    return createQuery(new DatabaseQuery(singletonList(table), new SimpleSQLiteQuery(sql, args)));
  }

  /**
   * See {@link #createQuery(String, String, Object...)} for usage. This overload allows for
   * monitoring multiple tables for changes.
   *
   * @see SupportSQLiteDatabase#query(String, Object[])
   */
  @CheckResult @NonNull
  public QueryObservable createQuery(@NonNull final Iterable<String> tables, @NonNull String sql,
      @NonNull Object... args) {
    return createQuery(new DatabaseQuery(tables, new SimpleSQLiteQuery(sql, args)));
  }

  /**
   * Create an observable which will notify subscribers with a {@linkplain Query query} for
   * execution. Subscribers are responsible for always closing {@link Cursor} instance
   * returned from the {@link Query}.
   *

   * Subscribers will receive an immediate notification for initial data as well as subsequent
   * notifications for when the supplied {@code table}'s data changes through the {@code insert},
   * {@code update}, and {@code delete} methods of this class. Unsubscribe when you no longer want
   * updates to a query.
   *

   * Since database triggers are inherently asynchronous, items emitted from the returned
   * observable use the {@link Scheduler} supplied to {@link SqlBrite#wrapDatabaseHelper}. For
   * consistency, the immediate notification sent on subscribe also uses this scheduler. As such,
   * calling {@link Observable#subscribeOn subscribeOn} on the returned observable has no effect.
   *

   * Note: To skip the immediate notification and only receive subsequent notifications when data
   * has changed call {@code skip(1)} on the returned observable.
   *

   * Warning: this method does not perform the query! Only by subscribing to the returned
   * {@link Observable} will the operation occur.
   *
   * @see SupportSQLiteDatabase#query(SupportSQLiteQuery)
   */
  @CheckResult @NonNull
  public QueryObservable createQuery(@NonNull final String table,
      @NonNull SupportSQLiteQuery query) {
    return createQuery(new DatabaseQuery(singletonList(table), query));
  }

  /**
   * See {@link #createQuery(String, SupportSQLiteQuery)} for usage. This overload allows for
   * monitoring multiple tables for changes.
   *
   * @see SupportSQLiteDatabase#query(SupportSQLiteQuery)
   */
  @CheckResult @NonNull
  public QueryObservable createQuery(@NonNull final Iterable<String> tables,
      @NonNull SupportSQLiteQuery query) {
    return createQuery(new DatabaseQuery(tables, query));
  }

  @CheckResult @NonNull
  private QueryObservable createQuery(DatabaseQuery query) {
    if (transactions.get() != null) {
      throw new IllegalStateException("Cannot create observable query in transaction. "
          + "Use query() for a query inside a transaction.");
    }

    return triggers //
        .filter(query) // DatabaseQuery filters triggers to on tables we care about.
        .map(query) // DatabaseQuery maps to itself to save an allocation.
        .startWith(query) //
        .observeOn(scheduler) //
        .compose(queryTransformer) // Apply the user's query transformer.
        .doOnSubscribe(ensureNotInTransaction)
        .to(QUERY_OBSERVABLE);
  }

  /**
   * Runs the provided SQL and returns a {@link Cursor} over the result set.
   *
   * @see SupportSQLiteDatabase#query(String, Object[])
   */
  @CheckResult @WorkerThread
  public Cursor query(@NonNull String sql, @NonNull Object... args) {
    Cursor cursor = getReadableDatabase().query(sql, args);
    if (logging) {
      log("QUERY\n  sql: %s\n  args: %s", indentSql(sql), Arrays.toString(args));
    }

    return cursor;
  }

  /**
   * Runs the provided {@link SupportSQLiteQuery} and returns a {@link Cursor} over the result set.
   *
   * @see SupportSQLiteDatabase#query(SupportSQLiteQuery)
   */
  @CheckResult @WorkerThread
  public Cursor query(@NonNull SupportSQLiteQuery query) {
    Cursor cursor = getReadableDatabase().query(query);
    if (logging) {
      log("QUERY\n  sql: %s", indentSql(query.getSql()));
    }

    return cursor;
  }

  /**
   * Insert a row into the specified {@code table} and notify any subscribed queries.
   *
   * @see SupportSQLiteDatabase#insert(String, int, ContentValues)
   */
  @WorkerThread
  public long insert(@NonNull String table, @ConflictAlgorithm int conflictAlgorithm,
      @NonNull ContentValues values) {
    SupportSQLiteDatabase db = getWritableDatabase();

    if (logging) {
      log("INSERT\n  table: %s\n  values: %s\n  conflictAlgorithm: %s", table, values,
          conflictString(conflictAlgorithm));
    }
    long rowId = db.insert(table, conflictAlgorithm, values);

    if (logging) log("INSERT id: %s", rowId);

    if (rowId != -1) {
      // Only send a table trigger if the insert was successful.
      sendTableTrigger(Collections.singleton(table));
    }
    return rowId;
  }

  /**
   * Delete rows from the specified {@code table} and notify any subscribed queries. This method
   * will not trigger a notification if no rows were deleted.
   *
   * @see SupportSQLiteDatabase#delete(String, String, Object[])
   */
  @WorkerThread
  public int delete(@NonNull String table, @Nullable String whereClause,
      @Nullable String... whereArgs) {
    SupportSQLiteDatabase db = getWritableDatabase();

    if (logging) {
      log("DELETE\n  table: %s\n  whereClause: %s\n  whereArgs: %s", table, whereClause,
          Arrays.toString(whereArgs));
    }
    int rows = db.delete(table, whereClause, whereArgs);

    if (logging) log("DELETE affected %s %s", rows, rows != 1 ? "rows" : "row");

    if (rows > 0) {
      // Only send a table trigger if rows were affected.
      sendTableTrigger(Collections.singleton(table));
    }
    return rows;
  }

  /**
   * Update rows in the specified {@code table} and notify any subscribed queries. This method
   * will not trigger a notification if no rows were updated.
   *
   * @see SupportSQLiteDatabase#update(String, int, ContentValues, String, Object[])
   */
  @WorkerThread
  public int update(@NonNull String table, @ConflictAlgorithm int conflictAlgorithm,
      @NonNull ContentValues values, @Nullable String whereClause, @Nullable String... whereArgs) {
    SupportSQLiteDatabase db = getWritableDatabase();

    if (logging) {
      log("UPDATE\n  table: %s\n  values: %s\n  whereClause: %s\n  whereArgs: %s\n  conflictAlgorithm: %s",
          table, values, whereClause, Arrays.toString(whereArgs),
          conflictString(conflictAlgorithm));
    }
    int rows = db.update(table, conflictAlgorithm, values, whereClause, whereArgs);

    if (logging) log("UPDATE affected %s %s", rows, rows != 1 ? "rows" : "row");

    if (rows > 0) {
      // Only send a table trigger if rows were affected.
      sendTableTrigger(Collections.singleton(table));
    }
    return rows;
  }

  /**
   * Execute {@code sql} provided it is NOT a {@code SELECT} or any other SQL statement that
   * returns data. No data can be returned (such as the number of affected rows). Instead, use
   * {@link #insert}, {@link #update}, et al, when possible.
   *

   * No notifications will be sent to queries if {@code sql} affects the data of a table.
   *
   * @see SupportSQLiteDatabase#execSQL(String)
   */
  @WorkerThread
  public void execute(String sql) {
    if (logging) log("EXECUTE\n  sql: %s", indentSql(sql));

    getWritableDatabase().execSQL(sql);
  }

  /**
   * Execute {@code sql} provided it is NOT a {@code SELECT} or any other SQL statement that
   * returns data. No data can be returned (such as the number of affected rows). Instead, use
   * {@link #insert}, {@link #update}, et al, when possible.
   *

   * No notifications will be sent to queries if {@code sql} affects the data of a table.
   *
   * @see SupportSQLiteDatabase#execSQL(String, Object[])
   */
  @WorkerThread
  public void execute(String sql, Object... args) {
    if (logging) log("EXECUTE\n  sql: %s\n  args: %s", indentSql(sql), Arrays.toString(args));

    getWritableDatabase().execSQL(sql, args);
  }

  /**
   * Execute {@code sql} provided it is NOT a {@code SELECT} or any other SQL statement that
   * returns data. No data can be returned (such as the number of affected rows). Instead, use
   * {@link #insert}, {@link #update}, et al, when possible.
   *

   * A notification to queries for {@code table} will be sent after the statement is executed.
   *
   * @see SupportSQLiteDatabase#execSQL(String)
   */
  @WorkerThread
  public void executeAndTrigger(String table, String sql) {
    executeAndTrigger(Collections.singleton(table), sql);
  }

  /**
   * See {@link #executeAndTrigger(String, String)} for usage. This overload allows for triggering multiple tables.
   *
   * @see BriteDatabase#executeAndTrigger(String, String)
   */
  @WorkerThread
  public void executeAndTrigger(Set<String> tables, String sql) {
    execute(sql);

    sendTableTrigger(tables);
  }

  /**
   * Execute {@code sql} provided it is NOT a {@code SELECT} or any other SQL statement that
   * returns data. No data can be returned (such as the number of affected rows). Instead, use
   * {@link #insert}, {@link #update}, et al, when possible.
   *

   * A notification to queries for {@code table} will be sent after the statement is executed.
   *
   * @see SupportSQLiteDatabase#execSQL(String, Object[])
   */
  @WorkerThread
  public void executeAndTrigger(String table, String sql, Object... args) {
    executeAndTrigger(Collections.singleton(table), sql, args);
  }

  /**
   * See {@link #executeAndTrigger(String, String, Object...)} for usage. This overload allows for triggering multiple tables.
   *
   * @see BriteDatabase#executeAndTrigger(String, String, Object...)
   */
  @WorkerThread
  public void executeAndTrigger(Set<String> tables, String sql, Object... args) {
    execute(sql, args);

    sendTableTrigger(tables);
  }

  /**
   * Execute {@code statement}, if the the number of rows affected by execution of this SQL
   * statement is of any importance to the caller - for example, UPDATE / DELETE SQL statements.
   *
   * @return the number of rows affected by this SQL statement execution.
   * @throws android.database.SQLException If the SQL string is invalid
   *
   * @see SupportSQLiteStatement#executeUpdateDelete()
   */
  @WorkerThread
  public int executeUpdateDelete(String table, SupportSQLiteStatement statement) {
    return executeUpdateDelete(Collections.singleton(table), statement);
  }

  /**
   * See {@link #executeUpdateDelete(String, SupportSQLiteStatement)} for usage. This overload
   * allows for triggering multiple tables.
   *
   * @see BriteDatabase#executeUpdateDelete(String, SupportSQLiteStatement)
   */
  @WorkerThread
  public int executeUpdateDelete(Set<String> tables, SupportSQLiteStatement statement) {
    if (logging) log("EXECUTE\n %s", statement);

    int rows = statement.executeUpdateDelete();
    if (rows > 0) {
      // Only send a table trigger if rows were affected.
      sendTableTrigger(tables);
    }
    return rows;
  }

  /**
   * Execute {@code statement} and return the ID of the row inserted due to this call.
   * The SQL statement should be an INSERT for this to be a useful call.
   *
   * @return the row ID of the last row inserted, if this insert is successful. -1 otherwise.
   *
   * @throws android.database.SQLException If the SQL string is invalid
   *
   * @see SupportSQLiteStatement#executeInsert()
   */
  @WorkerThread
  public long executeInsert(String table, SupportSQLiteStatement statement) {
    return executeInsert(Collections.singleton(table), statement);
  }

  /**
   * See {@link #executeInsert(String, SupportSQLiteStatement)} for usage. This overload allows for
   * triggering multiple tables.
   *
   * @see BriteDatabase#executeInsert(String, SupportSQLiteStatement)
   */
  @WorkerThread
  public long executeInsert(Set<String> tables, SupportSQLiteStatement statement) {
    if (logging) log("EXECUTE\n %s", statement);

    long rowId = statement.executeInsert();
    if (rowId != -1) {
      // Only send a table trigger if the insert was successful.
      sendTableTrigger(tables);
    }
    return rowId;
  }

  /** An in-progress database transaction. */
  public interface Transaction extends Closeable {
    /**
     * End a transaction. See {@link #newTransaction()} for notes about how to use this and when
     * transactions are committed and rolled back.
     *
     * @see SupportSQLiteDatabase#endTransaction()
     */
    @WorkerThread
    void end();

    /**
     * Marks the current transaction as successful. Do not do any more database work between
     * calling this and calling {@link #end()}. Do as little non-database work as possible in that
     * situation too. If any errors are encountered between this and {@link #end()} the transaction
     * will still be committed.
     *
     * @see SupportSQLiteDatabase#setTransactionSuccessful()
     */
    @WorkerThread
    void markSuccessful();

    /**
     * Temporarily end the transaction to let other threads run. The transaction is assumed to be
     * successful so far. Do not call {@link #markSuccessful()} before calling this. When this
     * returns a new transaction will have been created but not marked as successful. This assumes
     * that there are no nested transactions (newTransaction has only been called once) and will
     * throw an exception if that is not the case.
     *
     * @return true if the transaction was yielded
     *
     * @see SupportSQLiteDatabase#yieldIfContendedSafely()
     */
    @WorkerThread
    boolean yieldIfContendedSafely();

    /**
     * Temporarily end the transaction to let other threads run. The transaction is assumed to be
     * successful so far. Do not call {@link #markSuccessful()} before calling this. When this
     * returns a new transaction will have been created but not marked as successful. This assumes
     * that there are no nested transactions (newTransaction has only been called once) and will
     * throw an exception if that is not the case.
     *
     * @param sleepAmount if > 0, sleep this long before starting a new transaction if
     *   the lock was actually yielded. This will allow other background threads to make some
     *   more progress than they would if we started the transaction immediately.
     * @return true if the transaction was yielded
     *
     * @see SupportSQLiteDatabase#yieldIfContendedSafely(long)
     */
    @WorkerThread
    boolean yieldIfContendedSafely(long sleepAmount, TimeUnit sleepUnit);

    /**
     * Equivalent to calling {@link #end()}
     */
    @WorkerThread
    @Override void close();
  }

  @IntDef({
      CONFLICT_ABORT,
      CONFLICT_FAIL,
      CONFLICT_IGNORE,
      CONFLICT_NONE,
      CONFLICT_REPLACE,
      CONFLICT_ROLLBACK
  })
  @Retention(SOURCE)
  private @interface ConflictAlgorithm {
  }

  static String indentSql(String sql) {
    return sql.replace("\n", "\n       ");
  }

  void log(String message, Object... args) {
    if (args.length > 0) message = String.format(message, args);
    logger.log(message);
  }

  private static String conflictString(@ConflictAlgorithm int conflictAlgorithm) {
    switch (conflictAlgorithm) {
      case CONFLICT_ABORT:
        return "abort";
      case CONFLICT_FAIL:
        return "fail";
      case CONFLICT_IGNORE:
        return "ignore";
      case CONFLICT_NONE:
        return "none";
      case CONFLICT_REPLACE:
        return "replace";
      case CONFLICT_ROLLBACK:
        return "rollback";
      default:
        return "unknown (" + conflictAlgorithm + ')';
    }
  }

  static final class SqliteTransaction extends LinkedHashSet<String>
      implements SQLiteTransactionListener {
    final SqliteTransaction parent;
    boolean commit;

    SqliteTransaction(SqliteTransaction parent) {
      this.parent = parent;
    }

    @Override public void onBegin() {
    }

    @Override public void onCommit() {
      commit = true;
    }

    @Override public void onRollback() {
    }

    @Override public String toString() {
      String name = String.format("%08x", System.identityHashCode(this));
      return parent == null ? name : name + " [" + parent.toString() + ']';
    }
  }

  final class DatabaseQuery extends Query
      implements Function<Set<String>, Query>, Predicate<Set<String>> {
    private final Iterable<String> tables;
    private final SupportSQLiteQuery query;

    DatabaseQuery(Iterable<String> tables, SupportSQLiteQuery query) {
      this.tables = tables;
      this.query = query;
    }

    @Override public Cursor run() {
      if (transactions.get() != null) {
        throw new IllegalStateException("Cannot execute observable query in a transaction.");
      }

      Cursor cursor = getReadableDatabase().query(query);

      if (logging) {
        log("QUERY\n  tables: %s\n  sql: %s", tables, indentSql(query.getSql()));
      }

      return cursor;
    }

    @Override public String toString() {
      return query.getSql();
    }

    @Override public Query apply(Set<String> ignored) {
      return this;
    }

    @Override public boolean test(Set<String> strings) {
      for (String table : tables) {
        if (strings.contains(table)) {
          return true;
        }
      }
      return false;
    }
  }
}


/*
 * Copyright (C) 2015 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.squareup.sqlbrite3;

import android.arch.persistence.db.SupportSQLiteOpenHelper;
import android.content.ContentResolver;
import android.database.Cursor;
import android.os.Build;
import android.support.annotation.CheckResult;
import android.support.annotation.NonNull;
import android.support.annotation.Nullable;
import android.support.annotation.RequiresApi;
import android.support.annotation.WorkerThread;
import android.util.Log;
import io.reactivex.Observable;
import io.reactivex.ObservableEmitter;
import io.reactivex.ObservableOnSubscribe;
import io.reactivex.ObservableOperator;
import io.reactivex.ObservableTransformer;
import io.reactivex.Scheduler;
import io.reactivex.functions.Function;
import java.util.List;
import java.util.Optional;

/**
 * A lightweight wrapper around {@link SupportSQLiteOpenHelper} which allows for continuously
 * observing the result of a query.
 */
public final class SqlBrite {
  static final Logger DEFAULT_LOGGER = new Logger() {
    @Override public void log(String message) {
      Log.d("SqlBrite", message);
    }
  };
  static final ObservableTransformer<Query, Query> DEFAULT_TRANSFORMER =
      new ObservableTransformer<Query, Query>() {
        @Override public Observable<Query> apply(Observable<Query> queryObservable) {
          return queryObservable;
        }
      };

  public static final class Builder {
    private Logger logger = DEFAULT_LOGGER;
    private ObservableTransformer<Query, Query> queryTransformer = DEFAULT_TRANSFORMER;

    @CheckResult
    public Builder logger(@NonNull Logger logger) {
      if (logger == null) throw new NullPointerException("logger == null");
      this.logger = logger;
      return this;
    }

    @CheckResult
    public Builder queryTransformer(@NonNull ObservableTransformer<Query, Query> queryTransformer) {
      if (queryTransformer == null) throw new NullPointerException("queryTransformer == null");
      this.queryTransformer = queryTransformer;
      return this;
    }

    @CheckResult
    public SqlBrite build() {
      return new SqlBrite(logger, queryTransformer);
    }
  }

  final Logger logger;
  final ObservableTransformer<Query, Query> queryTransformer;

  SqlBrite(@NonNull Logger logger, @NonNull ObservableTransformer<Query, Query> queryTransformer) {
    this.logger = logger;
    this.queryTransformer = queryTransformer;
  }

  /**
   * Wrap a {@link SupportSQLiteOpenHelper} for observable queries.
   *

   * While not strictly required, instances of this class assume that they will be the only ones
   * interacting with the underlying {@link SupportSQLiteOpenHelper} and it is required for
   * automatic notifications of table changes to work. See {@linkplain BriteDatabase#createQuery the
   * query method} for more information on that behavior.
   *
   * @param scheduler The {@link Scheduler} on which items from {@link BriteDatabase#createQuery}
   * will be emitted.
   */
  @CheckResult @NonNull public BriteDatabase wrapDatabaseHelper(
      @NonNull SupportSQLiteOpenHelper helper,
      @NonNull Scheduler scheduler) {
    return new BriteDatabase(helper, logger, scheduler, queryTransformer);
  }

  /**
   * Wrap a {@link ContentResolver} for observable queries.
   *
   * @param scheduler The {@link Scheduler} on which items from
   * {@link BriteContentResolver#createQuery} will be emitted.
   */
  @CheckResult @NonNull public BriteContentResolver wrapContentProvider(
      @NonNull ContentResolver contentResolver, @NonNull Scheduler scheduler) {
    return new BriteContentResolver(contentResolver, logger, scheduler, queryTransformer);
  }

  /** An executable query. */
  public static abstract class Query {
    /**
     * Creates an {@linkplain ObservableOperator operator} which transforms a query returning a
     * single row to a {@code T} using {@code mapper}. Use with {@link Observable#lift}.
     *

     * It is an error for a query to pass through this operator with more than 1 row in its result
     * set. Use {@code LIMIT 1} on the underlying SQL query to prevent this. Result sets with 0 rows
     * do not emit an item.
     *

     * This operator ignores {@code null} cursors returned from {@link #run()}.
     *
     * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
     */
    @CheckResult @NonNull //
    public static <T> ObservableOperator<T, Query> mapToOne(@NonNull Function<Cursor, T> mapper) {
      return new QueryToOneOperator<>(mapper, null);
    }

    /**
     * Creates an {@linkplain ObservableOperator operator} which transforms a query returning a
     * single row to a {@code T} using {@code mapper}. Use with {@link Observable#lift}.
     *

     * It is an error for a query to pass through this operator with more than 1 row in its result
     * set. Use {@code LIMIT 1} on the underlying SQL query to prevent this. Result sets with 0 rows
     * emit {@code defaultValue}.
     *

     * This operator emits {@code defaultValue} if {@code null} is returned from {@link #run()}.
     *
     * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
     * @param defaultValue Value returned if result set is empty
     */
    @SuppressWarnings("ConstantConditions") // Public API contract.
    @CheckResult @NonNull
    public static <T> ObservableOperator<T, Query> mapToOneOrDefault(
        @NonNull Function<Cursor, T> mapper, @NonNull T defaultValue) {
      if (defaultValue == null) throw new NullPointerException("defaultValue == null");
      return new QueryToOneOperator<>(mapper, defaultValue);
    }

    /**
     * Creates an {@linkplain ObservableOperator operator} which transforms a query returning a
     * single row to a {@code Optional<T>} using {@code mapper}. Use with {@link Observable#lift}.
     *

     * It is an error for a query to pass through this operator with more than 1 row in its result
     * set. Use {@code LIMIT 1} on the underlying SQL query to prevent this. Result sets with 0 rows
     * emit {@link Optional#empty() Optional.empty()}.
     *

     * This operator ignores {@code null} cursors returned from {@link #run()}.
     *
     * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
     */
    @RequiresApi(Build.VERSION_CODES.N) //
    @CheckResult @NonNull //
    public static <T> ObservableOperator<Optional<T>, Query> mapToOptional(
        @NonNull Function<Cursor, T> mapper) {
      return new QueryToOptionalOperator<>(mapper);
    }

    /**
     * Creates an {@linkplain ObservableOperator operator} which transforms a query to a
     * {@code List<T>} using {@code mapper}. Use with {@link Observable#lift}.
     *

     * Be careful using this operator as it will always consume the entire cursor and create objects
     * for each row, every time this observable emits a new query. On tables whose queries update
     * frequently or very large result sets this can result in the creation of many objects.
     *

     * This operator ignores {@code null} cursors returned from {@link #run()}.
     *
     * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
     */
    @CheckResult @NonNull
    public static <T> ObservableOperator<List<T>, Query> mapToList(
        @NonNull Function<Cursor, T> mapper) {
      return new QueryToListOperator<>(mapper);
    }

    /**
     * Execute the query on the underlying database and return the resulting cursor.
     *
     * @return A {@link Cursor} with query results, or {@code null} when the query could not be
     * executed due to a problem with the underlying store. Unfortunately it is not well documented
     * when {@code null} is returned. It usually involves a problem in communicating with the
     * underlying store and should either be treated as failure or ignored for retry at a later
     * time.
     */
    @CheckResult @WorkerThread
    @Nullable
    public abstract Cursor run();

    /**
     * Execute the query on the underlying database and return an Observable of each row mapped to
     * {@code T} by {@code mapper}.
     *

     * Standard usage of this operation is in {@code flatMap}:
     *
{@code
     * flatMap(q -> q.asRows(Item.MAPPER).toList())
     * }

     * However, the above is a more-verbose but identical operation as
     * {@link QueryObservable#mapToList}. This {@code asRows} method should be used when you need
     * to limit or filter the items separate from the actual query.
     *
{@code
     * flatMap(q -> q.asRows(Item.MAPPER).take(5).toList())
     * // or...
     * flatMap(q -> q.asRows(Item.MAPPER).filter(i -> i.isActive).toList())
     * }

     *

     * Note: Limiting results or filtering will almost always be faster in the database as part of
     * a query and should be preferred, where possible.
     *

     * The resulting observable will be empty if {@code null} is returned from {@link #run()}.
     */
    @CheckResult @NonNull
    public final <T> Observable<T> asRows(final Function<Cursor, T> mapper) {
      return Observable.create(new ObservableOnSubscribe<T>() {
        @Override public void subscribe(ObservableEmitter<T> e) throws Exception {
          Cursor cursor = run();
          if (cursor != null) {
            try {
              while (cursor.moveToNext() && !e.isDisposed()) {
                e.onNext(mapper.apply(cursor));
              }
            } finally {
              cursor.close();
            }
          }
          if (!e.isDisposed()) {
            e.onComplete();
          }
        }
      });
    }
  }

  /** A simple indirection for logging debug messages. */
  public interface Logger {
    void log(String message);
  }
}


/*
 * Copyright (C) 2015 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.squareup.sqlbrite3;

import android.content.ContentResolver;
import android.database.ContentObserver;
import android.database.Cursor;
import android.net.Uri;
import android.os.Handler;
import android.os.Looper;
import android.support.annotation.CheckResult;
import android.support.annotation.NonNull;
import android.support.annotation.Nullable;
import com.squareup.sqlbrite3.SqlBrite.Logger;
import com.squareup.sqlbrite3.SqlBrite.Query;
import io.reactivex.Observable;
import io.reactivex.ObservableEmitter;
import io.reactivex.ObservableOnSubscribe;
import io.reactivex.ObservableTransformer;
import io.reactivex.Scheduler;
import io.reactivex.functions.Cancellable;
import java.util.Arrays;

import static com.squareup.sqlbrite3.QueryObservable.QUERY_OBSERVABLE;
import static java.lang.System.nanoTime;
import static java.util.concurrent.TimeUnit.NANOSECONDS;

/**
 * A lightweight wrapper around {@link ContentResolver} which allows for continuously observing
 * the result of a query. Create using a {@link SqlBrite} instance.
 */
public final class BriteContentResolver {
  final Handler contentObserverHandler = new Handler(Looper.getMainLooper());

  final ContentResolver contentResolver;
  private final Logger logger;
  private final Scheduler scheduler;
  private final ObservableTransformer<Query, Query> queryTransformer;

  volatile boolean logging;

  BriteContentResolver(ContentResolver contentResolver, Logger logger, Scheduler scheduler,
      ObservableTransformer<Query, Query> queryTransformer) {
    this.contentResolver = contentResolver;
    this.logger = logger;
    this.scheduler = scheduler;
    this.queryTransformer = queryTransformer;
  }

  /** Control whether debug logging is enabled. */
  public void setLoggingEnabled(boolean enabled) {
    logging = enabled;
  }

  /**
   * Create an observable which will notify subscribers with a {@linkplain Query query} for
   * execution. Subscribers are responsible for always closing {@link Cursor} instance
   * returned from the {@link Query}.
   *

   * Subscribers will receive an immediate notification for initial data as well as subsequent
   * notifications for when the supplied {@code uri}'s data changes. Unsubscribe when you no longer
   * want updates to a query.
   *

   * Since content resolver triggers are inherently asynchronous, items emitted from the returned
   * observable use the {@link Scheduler} supplied to {@link SqlBrite#wrapContentProvider}. For
   * consistency, the immediate notification sent on subscribe also uses this scheduler. As such,
   * calling {@link Observable#subscribeOn subscribeOn} on the returned observable has no effect.
   *

   * Note: To skip the immediate notification and only receive subsequent notifications when data
   * has changed call {@code skip(1)} on the returned observable.
   *

   * Warning: this method does not perform the query! Only by subscribing to the returned
   * {@link Observable} will the operation occur.
   *
   * @see ContentResolver#query(Uri, String[], String, String[], String)
   * @see ContentResolver#registerContentObserver(Uri, boolean, ContentObserver)
   */
  @CheckResult @NonNull
  public QueryObservable createQuery(@NonNull final Uri uri, @Nullable final String[] projection,
      @Nullable final String selection, @Nullable final String[] selectionArgs, @Nullable
      final String sortOrder, final boolean notifyForDescendents) {
    final Query query = new Query() {
      @Override public Cursor run() {
        long startNanos = nanoTime();
        Cursor cursor = contentResolver.query(uri, projection, selection, selectionArgs, sortOrder);

        if (logging) {
          long tookMillis = NANOSECONDS.toMillis(nanoTime() - startNanos);
          log("QUERY (%sms)\n  uri: %s\n  projection: %s\n  selection: %s\n  selectionArgs: %s\n  "
                  + "sortOrder: %s\n  notifyForDescendents: %s", tookMillis, uri,
              Arrays.toString(projection), selection, Arrays.toString(selectionArgs), sortOrder,
              notifyForDescendents);
        }

        return cursor;
      }
    };
    Observable<Query> queries = Observable.create(new ObservableOnSubscribe<Query>() {
      @Override public void subscribe(final ObservableEmitter<Query> e) throws Exception {
        final ContentObserver observer = new ContentObserver(contentObserverHandler) {
          @Override public void onChange(boolean selfChange) {
            if (!e.isDisposed()) {
              e.onNext(query);
            }
          }
        };
        contentResolver.registerContentObserver(uri, notifyForDescendents, observer);
        e.setCancellable(new Cancellable() {
          @Override public void cancel() throws Exception {
            contentResolver.unregisterContentObserver(observer);
          }
        });

        if (!e.isDisposed()) {
          e.onNext(query); // Trigger initial query.
        }
      }
    });
    return queries //
        .observeOn(scheduler) //
        .compose(queryTransformer) // Apply the user's query transformer.
        .to(QUERY_OBSERVABLE);
  }

  void log(String message, Object... args) {
    if (args.length > 0) message = String.format(message, args);
    logger.log(message);
  }
}


package com.squareup.sqlbrite3;

import android.database.Cursor;
import android.os.Build;
import android.support.annotation.CheckResult;
import android.support.annotation.NonNull;
import android.support.annotation.RequiresApi;
import com.squareup.sqlbrite3.SqlBrite.Query;
import io.reactivex.Observable;
import io.reactivex.Observer;
import io.reactivex.functions.Function;
import java.util.List;
import java.util.Optional;

/** An {@link Observable} of {@link Query} which offers query-specific convenience operators. */
public final class QueryObservable extends Observable<Query> {
  static final Function<Observable<Query>, QueryObservable> QUERY_OBSERVABLE =
      new Function<Observable<Query>, QueryObservable>() {
        @Override public QueryObservable apply(Observable<Query> queryObservable) {
          return new QueryObservable(queryObservable);
        }
      };

  private final Observable<Query> upstream;

  public QueryObservable(Observable<Query> upstream) {
    this.upstream = upstream;
  }

  @Override protected void subscribeActual(Observer<? super Query> observer) {
    upstream.subscribe(observer);
  }

  /**
   * Given a function mapping the current row of a {@link Cursor} to {@code T}, transform each
   * emitted {@link Query} which returns a single row to {@code T}.
   *

   * It is an error for a query to pass through this operator with more than 1 row in its result
   * set. Use {@code LIMIT 1} on the underlying SQL query to prevent this. Result sets with 0 rows
   * do not emit an item.
   *

   * This method is equivalent to:
   *
{@code
   * flatMap(q -> q.asRows(mapper).take(1))
   * }

   * and a convenience operator for:
   *
{@code
   * lift(Query.mapToOne(mapper))
   * }

   *
   * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
   */
  @CheckResult @NonNull
  public final <T> Observable<T> mapToOne(@NonNull Function<Cursor, T> mapper) {
    return lift(Query.mapToOne(mapper));
  }

  /**
   * Given a function mapping the current row of a {@link Cursor} to {@code T}, transform each
   * emitted {@link Query} which returns a single row to {@code T}.
   *

   * It is an error for a query to pass through this operator with more than 1 row in its result
   * set. Use {@code LIMIT 1} on the underlying SQL query to prevent this. Result sets with 0 rows
   * emit {@code defaultValue}.
   *

   * This method is equivalent to:
   *
{@code
   * flatMap(q -> q.asRows(mapper).take(1).defaultIfEmpty(defaultValue))
   * }

   * and a convenience operator for:
   *
{@code
   * lift(Query.mapToOneOrDefault(mapper, defaultValue))
   * }

   *
   * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
   * @param defaultValue Value returned if result set is empty
   */
  @CheckResult @NonNull
  public final <T> Observable<T> mapToOneOrDefault(@NonNull Function<Cursor, T> mapper,
      @NonNull T defaultValue) {
    return lift(Query.mapToOneOrDefault(mapper, defaultValue));
  }

  /**
   * Given a function mapping the current row of a {@link Cursor} to {@code T}, transform each
   * emitted {@link Query} which returns a single row to {@code Optional<T>}.
   *

   * It is an error for a query to pass through this operator with more than 1 row in its result
   * set. Use {@code LIMIT 1} on the underlying SQL query to prevent this. Result sets with 0 rows
   * emit {@link Optional#empty() Optional.empty()}
   *

   * This method is equivalent to:
   *
{@code
   * flatMap(q -> q.asRows(mapper).take(1).map(Optional::of).defaultIfEmpty(Optional.empty())
   * }

   * and a convenience operator for:
   *
{@code
   * lift(Query.mapToOptional(mapper))
   * }

   *
   * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
   */
  @RequiresApi(Build.VERSION_CODES.N)
  @CheckResult @NonNull
  public final <T> Observable<Optional<T>> mapToOptional(@NonNull Function<Cursor, T> mapper) {
    return lift(Query.mapToOptional(mapper));
  }

  /**
   * Given a function mapping the current row of a {@link Cursor} to {@code T}, transform each
   * emitted {@link Query} to a {@code List<T>}.
   *

   * Be careful using this operator as it will always consume the entire cursor and create objects
   * for each row, every time this observable emits a new query. On tables whose queries update
   * frequently or very large result sets this can result in the creation of many objects.
   *

   * This method is equivalent to:
   *
{@code
   * flatMap(q -> q.asRows(mapper).toList())
   * }

   * and a convenience operator for:
   *
{@code
   * lift(Query.mapToList(mapper))
   * }

   *

   * Consider using {@link Query#asRows} if you need to limit or filter in memory.
   *
   * @param mapper Maps the current {@link Cursor} row to {@code T}. May not return null.
   */
  @CheckResult @NonNull
  public final <T> Observable<List<T>> mapToList(@NonNull Function<Cursor, T> mapper) {
    return lift(Query.mapToList(mapper));
  }
}


/*
 * Copyright (c) 2014, Cloudera, Inc. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.app.rdf;

import java.util.BitSet;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import com.google.common.base.Preconditions;
import org.dmg.pmml.DataDictionary;
import org.dmg.pmml.MiningField;
import org.dmg.pmml.MiningFunction;
import org.dmg.pmml.MiningSchema;
import org.dmg.pmml.Model;
import org.dmg.pmml.PMML;
import org.dmg.pmml.Predicate;
import org.dmg.pmml.ScoreDistribution;
import org.dmg.pmml.SimplePredicate;
import org.dmg.pmml.SimpleSetPredicate;
import org.dmg.pmml.True;
import org.dmg.pmml.mining.MiningModel;
import org.dmg.pmml.mining.Segment;
import org.dmg.pmml.mining.Segmentation;
import org.dmg.pmml.tree.Node;
import org.dmg.pmml.tree.TreeModel;

import com.cloudera.oryx.app.classreg.predict.CategoricalPrediction;
import com.cloudera.oryx.app.classreg.predict.NumericPrediction;
import com.cloudera.oryx.app.classreg.predict.Prediction;
import com.cloudera.oryx.app.pmml.AppPMMLUtils;
import com.cloudera.oryx.app.rdf.decision.CategoricalDecision;
import com.cloudera.oryx.app.rdf.decision.Decision;
import com.cloudera.oryx.app.rdf.decision.NumericDecision;
import com.cloudera.oryx.app.rdf.tree.DecisionForest;
import com.cloudera.oryx.app.rdf.tree.DecisionNode;
import com.cloudera.oryx.app.rdf.tree.DecisionTree;
import com.cloudera.oryx.app.rdf.tree.TerminalNode;
import com.cloudera.oryx.app.rdf.tree.TreeNode;
import com.cloudera.oryx.app.schema.CategoricalValueEncodings;
import com.cloudera.oryx.app.schema.InputSchema;
import com.cloudera.oryx.common.collection.Pair;
import com.cloudera.oryx.common.text.TextUtils;

/**
 * Utility methods related to PMML serialization of random decision forest classification
 * and regression models.
 */
public final class RDFPMMLUtils {

  private RDFPMMLUtils() {}

  /**
   * Validates that the encoded PMML model received matches expected schema.
   *
   * @param pmml {@link PMML} encoding of a decision forest
   * @param schema expected schema attributes of decision forest
   */
  public static void validatePMMLVsSchema(PMML pmml, InputSchema schema) {
    List<Model> models = pmml.getModels();
    Preconditions.checkArgument(models.size() == 1,
                                "Should have exactly one model, but had %s", models.size());

    Model model = models.get(0);
    MiningFunction function = model.getMiningFunction();
    if (schema.isClassification()) {
      Preconditions.checkArgument(function == MiningFunction.CLASSIFICATION,
                                  "Expected classification function type but got %s",
                                  function);
    } else {
      Preconditions.checkArgument(function == MiningFunction.REGRESSION,
                                  "Expected regression function type but got %s",
                                  function);
    }

    DataDictionary dictionary = pmml.getDataDictionary();
    Preconditions.checkArgument(
        schema.getFeatureNames().equals(AppPMMLUtils.getFeatureNames(dictionary)),
        "Feature names in schema don't match names in PMML");

    MiningSchema miningSchema = model.getMiningSchema();
    Preconditions.checkArgument(schema.getFeatureNames().equals(
        AppPMMLUtils.getFeatureNames(miningSchema)));

    Integer pmmlIndex = AppPMMLUtils.findTargetIndex(miningSchema);
    if (schema.hasTarget()) {
      int schemaIndex = schema.getTargetFeatureIndex();
      Preconditions.checkArgument(
          pmmlIndex != null && schemaIndex == pmmlIndex,
          "Configured schema expects target at index %s, but PMML has target at index %s",
          schemaIndex, pmmlIndex);
    } else {
      Preconditions.checkArgument(pmmlIndex == null);
    }
  }

  /**
   * @param pmml PMML representation of decision forest
   * @return a {@link DecisionForest} representation of the PMML encoded model
   */
  public static Pair<DecisionForest,CategoricalValueEncodings> read(PMML pmml) {

    DataDictionary dictionary = pmml.getDataDictionary();
    List<String> featureNames = AppPMMLUtils.getFeatureNames(dictionary);
    CategoricalValueEncodings categoricalValueEncodings =
        AppPMMLUtils.buildCategoricalValueEncodings(dictionary);

    List<Model> models = pmml.getModels();
    Model model = models.get(0);
    MiningSchema miningSchema = model.getMiningSchema();
    int targetIndex = Objects.requireNonNull(AppPMMLUtils.findTargetIndex(miningSchema));

    DecisionTree[] trees;
    double[] weights;
    if (model instanceof MiningModel) {
      MiningModel miningModel = (MiningModel) model;
      Segmentation segmentation = miningModel.getSegmentation();
      Preconditions.checkArgument(
          segmentation.getMultipleModelMethod() == Segmentation.MultipleModelMethod.WEIGHTED_AVERAGE ||
          segmentation.getMultipleModelMethod() == Segmentation.MultipleModelMethod.WEIGHTED_MAJORITY_VOTE);
      List<Segment> segments = segmentation.getSegments();
      Preconditions.checkArgument(!segments.isEmpty());

      trees = new DecisionTree[segments.size()];
      weights = new double[trees.length];
      for (int i = 0; i < trees.length; i++) {
        Segment segment = segments.get(i);
        Preconditions.checkArgument(segment.getPredicate() instanceof True);
        weights[i] = segment.getWeight();
        TreeModel treeModel = (TreeModel) segment.getModel();
        TreeNode root = translateFromPMML(treeModel.getNode(),
                                          categoricalValueEncodings,
                                          featureNames,
                                          targetIndex);
        trees[i] = new DecisionTree(root);
      }
    } else {
      // Single tree model
      TreeNode root = translateFromPMML(((TreeModel) model).getNode(),
                                        categoricalValueEncodings,
                                        featureNames,
                                        targetIndex);
      trees = new DecisionTree[] { new DecisionTree(root) };
      weights = new double[] { 1.0 };
    }

    List<MiningField> miningFields = miningSchema.getMiningFields();
    double[] featureImportances = new double[featureNames.size()];
    for (int i = 0; i < miningFields.size(); i++) {
      MiningField field = miningFields.get(i);
      Double importance = field.getImportance();
      if (importance != null) {
        featureImportances[i] = importance;
      }
    }

    return new Pair<>(new DecisionForest(trees, weights, featureImportances),
                      categoricalValueEncodings);
  }

  private static TreeNode translateFromPMML(Node root,
                                            CategoricalValueEncodings categoricalValueEncodings,
                                            List<String> featureNames,
                                            int targetIndex) {

    String id = root.getId();
    List<Node> children = root.getNodes();
    if (children.isEmpty()) {
      // Terminal
      Collection<ScoreDistribution> scoreDistributions = root.getScoreDistributions();
      Prediction prediction;
      if (scoreDistributions != null && !scoreDistributions.isEmpty()) {
        // Categorical target
        Map<String,Integer> targetEncoding =
            categoricalValueEncodings.getValueEncodingMap(targetIndex);
        double[] categoryCounts = new double[targetEncoding.size()];
        for (ScoreDistribution dist : scoreDistributions) {
          int encoding = targetEncoding.get(dist.getValue());
          categoryCounts[encoding] = dist.getRecordCount();
        }
        prediction = new CategoricalPrediction(categoryCounts);
      } else {
        prediction = new NumericPrediction(Double.parseDouble(root.getScore()),
                                           (int) Math.round(root.getRecordCount()));
      }
      return new TerminalNode(id, prediction);
    }

    Preconditions.checkArgument(children.size() == 2);
    // Decision
    Node child1 = children.get(0);
    Node child2 = children.get(1);
    Node negativeLeftChild;
    Node positiveRightChild;
    if (child1.getPredicate() instanceof True) {
      negativeLeftChild = child1;
      positiveRightChild = child2;
    } else {
      Preconditions.checkArgument(child2.getPredicate() instanceof True);
      negativeLeftChild = child2;
      positiveRightChild = child1;
    }

    Decision decision;
    Predicate predicate = positiveRightChild.getPredicate();
    boolean defaultDecision = positiveRightChild.getId().equals(root.getDefaultChild());

    if (predicate instanceof SimplePredicate) {
      // Numeric decision
      SimplePredicate simplePredicate = (SimplePredicate) predicate;
      SimplePredicate.Operator operator = simplePredicate.getOperator();
      Preconditions.checkArgument(
          operator == SimplePredicate.Operator.GREATER_OR_EQUAL ||
          operator == SimplePredicate.Operator.GREATER_THAN);
      double threshold = Double.parseDouble(simplePredicate.getValue());
      // NumericDecision uses >= criteria. Increase threshold by one ulp to implement
      // "> threshold" as ">= (threshold + ulp)"
      if (operator == SimplePredicate.Operator.GREATER_THAN) {
        threshold += Math.ulp(threshold);
      }
      int featureNumber = featureNames.indexOf(simplePredicate.getField().getValue());
      decision = new NumericDecision(featureNumber, threshold, defaultDecision);

    } else {
      // Categorical decision
      Preconditions.checkArgument(predicate instanceof SimpleSetPredicate);
      SimpleSetPredicate simpleSetPredicate = (SimpleSetPredicate) predicate;
      SimpleSetPredicate.BooleanOperator operator = simpleSetPredicate.getBooleanOperator();
      Preconditions.checkArgument(
          operator == SimpleSetPredicate.BooleanOperator.IS_IN ||
              operator == SimpleSetPredicate.BooleanOperator.IS_NOT_IN);
      int featureNumber = featureNames.indexOf(simpleSetPredicate.getField().getValue());
      Map<String,Integer> valueEncodingMap =
          categoricalValueEncodings.getValueEncodingMap(featureNumber);
      String[] categories = TextUtils.parseDelimited(simpleSetPredicate.getArray().getValue(), ' ');
      BitSet activeCategories = new BitSet(valueEncodingMap.size());
      if (operator == SimpleSetPredicate.BooleanOperator.IS_IN) {
        for (String category : categories) {
          activeCategories.set(valueEncodingMap.get(category));
        }
      } else {
        // "not in"
        valueEncodingMap.values().forEach(activeCategories::set);
        for (String category : categories) {
          activeCategories.clear(valueEncodingMap.get(category));
        }
      }
      decision = new CategoricalDecision(featureNumber, activeCategories, defaultDecision);
    }

    return new DecisionNode(
        id,
        decision,
        translateFromPMML(negativeLeftChild,
                          categoricalValueEncodings,
                          featureNames,
                          targetIndex),
        translateFromPMML(positiveRightChild,
                          categoricalValueEncodings,
                          featureNames,
                          targetIndex));
  }

}


/*
 * Copyright (c) 2015, Cloudera, Inc. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.app.als;

import com.cloudera.oryx.common.math.Solver;
import com.cloudera.oryx.common.math.VectorMath;

/**
 * ALS-related utility methods for the app tier.
 */
public final class ALSUtils {

  private ALSUtils() {}

  /**
   * Computes how the estimated strength of interaction in the model should change -- to what target
   * value -- in response to a new interaction.
   *
   * @param implicit whether the model is implicit feedback
   * @param value new interaction's strength
   * @param currentValue existing estimated of strength of interaction
   * @return new target estimated of strength of interaction, or NaN to signal "no change needed"
   */
  public static double computeTargetQui(boolean implicit, double value, double currentValue) {
    // We want Qui to change based on value. What's the target value, Qui'?
    if (implicit) {
      // Target is really 1, or 0, depending on whether value is positive or negative.
      // This wouldn't account for the strength though. Instead the target is a function
      // of the current value and strength. If the current value is c, and value is positive
      // then the target is somewhere between c and 1 depending on the strength. If current
      // value is already >= 1, there's no effect. Similarly for negative values.
      if (value > 0.0f && currentValue < 1.0) {
        double diff = 1.0 - Math.max(0.0, currentValue);
        return currentValue + (value / (1.0 + value)) * diff;
      }
      if (value < 0.0f && currentValue > 0.0) {
        double diff = -Math.min(1.0, currentValue);
        return currentValue + (value / (value - 1.0)) * diff;
      }
      // No change
      return Double.NaN;
    } else {
      // Non-implicit -- value is supposed to be the new value
      return value;
    }
  }

  /**
   * Computes how a user vector Xu changes in response to interaction with an item vector Yi.
   * This can also be used to compute how an item vector changes in response to a user interaction,
   * even though the code naming follows the former convention.
   *
   * @param solver solver helping solve for Xu in Qu*Y = Xu * (Yt * Y)
   * @param value strength of interaction
   * @param Xu current user vector (null if no existing user vector)
   * @param Yi current item vector
   * @param implicit whether the model is implicit feedback
   * @return new user vector Xu, or {@code null} if no update should be made (i.e. there was no
   *  item vector; the update would push the new Qui farther out of range)
   */
  public static float[] computeUpdatedXu(Solver solver,
                                         double value,
                                         float[] Xu,
                                         float[] Yi,
                                         boolean implicit) {
    if (Yi == null) {
      return null;
    }

    boolean noXu = Xu == null;
    double Qui = noXu ? 0.0 : VectorMath.dot(Xu, Yi);
    // Qui' is the target, new value of Qui
    // 0.5 reflects a "don't know" state
    double targetQui = computeTargetQui(implicit, value, noXu ? 0.5 : Qui);
    if (Double.isNaN(targetQui)) {
      return null;
    }

    double dQui = targetQui - Qui;
    double[] dQuiYi = new double[Yi.length];
    for (int i = 0; i < dQuiYi.length; i++) {
      dQuiYi[i] = Yi[i] * dQui;
    }
    double[] dXu = solver.solveDToD(dQuiYi);

    float[] newXu = noXu ? new float[dXu.length] : Xu.clone();
    for (int i = 0; i < newXu.length; i++) {
      newXu[i] += dXu[i];
    }
    return newXu;
  }

}


/*
 * Copyright (c) 2015, Cloudera and Intel, Inc. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.app.kmeans;

import java.util.List;
import java.util.stream.Collectors;

import com.google.common.base.Preconditions;
import org.dmg.pmml.DataDictionary;
import org.dmg.pmml.MiningFunction;
import org.dmg.pmml.MiningSchema;
import org.dmg.pmml.Model;
import org.dmg.pmml.PMML;
import org.dmg.pmml.clustering.ClusteringModel;

import com.cloudera.oryx.app.pmml.AppPMMLUtils;
import com.cloudera.oryx.app.schema.InputSchema;
import com.cloudera.oryx.common.math.VectorMath;
import com.cloudera.oryx.common.text.TextUtils;

/**
 * Utility methods related to PMML serialization of k-means clustering models.
 */
public final class KMeansPMMLUtils {

  private KMeansPMMLUtils() {}

  /**
   * Validates that the encoded PMML model received matches expected schema.
   *
   * @param pmml {@link PMML} encoding of KMeans Clustering
   * @param schema expected schema attributes of KMeans Clustering
   */
  public static void validatePMMLVsSchema(PMML pmml, InputSchema schema) {
    List<Model> models = pmml.getModels();
    Preconditions.checkArgument(models.size() == 1,
        "Should have exactly one model, but had %s", models.size());

    Model model = models.get(0);
    Preconditions.checkArgument(model instanceof ClusteringModel);
    Preconditions.checkArgument(model.getMiningFunction() == MiningFunction.CLUSTERING);

    DataDictionary dictionary = pmml.getDataDictionary();
    Preconditions.checkArgument(
        schema.getFeatureNames().equals(AppPMMLUtils.getFeatureNames(dictionary)),
        "Feature names in schema don't match names in PMML");

    MiningSchema miningSchema = model.getMiningSchema();
    Preconditions.checkArgument(schema.getFeatureNames().equals(
        AppPMMLUtils.getFeatureNames(miningSchema)));

  }

  /**
   * @param pmml PMML representation of Clusters
   * @return List of {@link ClusterInfo}
   */
  public static List<ClusterInfo> read(PMML pmml) {
    Model model = pmml.getModels().get(0);
    Preconditions.checkArgument(model instanceof ClusteringModel);
    ClusteringModel clusteringModel = (ClusteringModel) model;

    return clusteringModel.getClusters().stream().map(cluster ->
      new ClusterInfo(Integer.parseInt(cluster.getId()),
                      VectorMath.parseVector(TextUtils.parseDelimited(cluster.getArray().getValue(), ' ')),
                      cluster.getSize())
    ).collect(Collectors.toList());
  }

}


/*
 * Copyright (c) 2014, Cloudera, Inc. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.common.math;

import java.util.Collection;

import com.github.fommil.netlib.BLAS;
import org.apache.commons.math3.random.RandomGenerator;

/**
 * Utility class with simple vector-related operations.
 */
public final class VectorMath {

  private static final BLAS blas = BLAS.getInstance();

  private VectorMath() {}

  /**
   * @return dot product of the two given arrays
   * @param x one array
   * @param y the other array
   */
  public static double dot(float[] x, float[] y) {
    int length = x.length;
    double dot = 0.0;
    for (int i = 0; i < length; i++) {
      dot += (double) x[i] * y[i];
    }
    return dot;
  }

  /**
   * @param x vector for whom norm to be calculated
   * @return the L2 norm of vector x
   */
  public static double norm(float[] x) {
    double total = 0.0;
    for (float f : x) {
      total += (double) f * f;
    }
    return Math.sqrt(total);
  }

  /**
   * @param x vector for whom norm to be calculated
   * @return the L2 norm of vector x
   */
  public static double norm(double[] x) {
    double total = 0.0;
    for (double d : x) {
      total += d * d;
    }
    return Math.sqrt(total);
  }

  /**
   * Computes cosine similarity of values in two given arrays, when the norm of one array is
   * known in advance, which is a not-uncommon case.
   *
   * @param x one array
   * @param y the other array
   * @param normY norm of y
   * @return cosine similarity = dot(x,y) / (norm(x) * norm(y))
   */
  public static double cosineSimilarity(float[] x, float[] y, double normY) {
    int length = x.length;
    double dot = 0.0;
    double totalXSq = 0.0;
    for (int i = 0; i < length; i++) {
      double xi = x[i];
      totalXSq += xi * xi;
      dot += xi * y[i];
    }
    return dot / (Math.sqrt(totalXSq) * normY);
  }

  /**
   * @param M tall, skinny matrix
   * @return MT * M as a dense lower-triangular matrix, represented in packed row-major form.
   */
  public static double[] transposeTimesSelf(Collection<float[]> M) {
    if (M == null || M.isEmpty()) {
      return null;
    }
    int features = M.iterator().next().length;
    double[] result = new double[features * (features + 1) / 2];
    double[] vectorD = new double[features];
    for (float[] vector : M) {
      // Unfortunately need to copy into a double[]
      for (int i = 0; i < vector.length; i++) {
        vectorD[i] = vector[i];
      }
      blas.dspr("L", features, 1.0, vectorD, 1, result);
    }
    return result;
  }

  /**
   * @param values numeric values as {@link String}s
   * @return values parsed as {@code double[]}
   */
  public static double[] parseVector(String[] values) {
    double[] doubles = new double[values.length];
    for (int i = 0; i < values.length; i++) {
      doubles[i] = Double.parseDouble(values[i]);
    }
    return doubles;
  }

  /**
   * @param features dimension of vector
   * @param random random number generator
   * @return vector whose direction from the origin is chosen uniformly at random, but which is not normalized
   */
  public static float[] randomVectorF(int features, RandomGenerator random) {
    float[] vector = new float[features];
    for (int i = 0; i < features; i++) {
      vector[i] = (float) random.nextGaussian();
    }
    return vector;
  }

}


/*
 * Copyright (c) 2016, Cloudera, Inc. and Intel Corp. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.app.als;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.cloudera.oryx.common.lang.LoggingCallable;
import com.cloudera.oryx.common.math.LinearSystemSolver;
import com.cloudera.oryx.common.math.Solver;

/**
 * A cache around a {@link Solver} that manages computing it when necessary from
 * {@link FeatureVectors}, tracking when it must be recomputed, and doing so
 * correctly in the presence of multiple threads.
 */
public final class SolverCache {

  private static final Logger log = LoggerFactory.getLogger(SolverCache.class);

  private final AtomicReference<Solver> solver;
  private final AtomicBoolean solverDirty;
  private final AtomicBoolean solverUpdating;
  private final CountDownLatch solverInitialized;
  private final Executor executor;
  private final FeatureVectors vectorPartitions;

  /**
   * @param executor {@link Executor} which should be used to asynchronously compute
   *  a {@link Solver}. Important: this should be able to execute more than 1 task in parallel.
   * @param vectorPartitions underling {@link FeatureVectors} data from which it should be
   *  computed
   */
  public SolverCache(Executor executor, FeatureVectors vectorPartitions) {
    solver = new AtomicReference<>();
    solverDirty = new AtomicBoolean(true);
    solverUpdating = new AtomicBoolean(false);
    solverInitialized = new CountDownLatch(1);
    this.executor = executor;
    this.vectorPartitions = vectorPartitions;
  }

  /**
   * Indicates that the underlying {@link FeatureVectors} have changed and so the
   * resulting {@link Solver} should be recomputed.
   */
  public void setDirty() {
    solverDirty.set(true);
  }

  /**
   * Proactively try to compute the solver asynchronously, if not already computed.
   * Does not block.
   */
  public void compute() {
    // Make sure only one attempts to build at one time
    if (solverUpdating.compareAndSet(false, true)) {
      executor.execute(LoggingCallable.log(() -> {
        try {
          log.info("Computing cached solver");
          // Not as much hurry if one already exists
          boolean lowPriority = solver.get() != null;
          Solver newYTYSolver = LinearSystemSolver.getSolver(vectorPartitions.getVTV(lowPriority));
          if (newYTYSolver != null) {
            log.info("Computed new solver {}", newYTYSolver);
            solver.set(newYTYSolver);
          }
        } finally {
          // Allow any threads waiting for initial model to proceed.
          // It's possible the solver is still null here if there is no input.
          solverInitialized.countDown();
          solverUpdating.set(false);
        }
      }).asRunnable());
    }
  }

  /**
   * @param blocking whether to block waiting for a first model
   * @return a recent {@link Solver}, optionally blocking if necessary to wait for an initial one to
   *  be computed. It does not block otherwise and returns the most recently computed one.
   *  Note that this method may return {@code null} even when blocking , for instance,
   *  if no solver is computable because there is no data.
   */
  public Solver get(boolean blocking) {
    if (solverDirty.getAndSet(false)) {
      // launch asynchronous update
      compute();
    }
    // Wait, in the case that there is no existing model already.
    // Otherwise this immediately proceeds.
    if (blocking && solverInitialized.getCount() > 0) { // Always 0 or 1
      // OK if countDown() happens here; await() will return immediately
      try {
        solverInitialized.await();
      } catch (InterruptedException e) {
        log.warn("Interrupted while waiting for model", e);
      }
    }
    return solver.get();
  }

}


/*
 * Copyright (c) 2015, Cloudera, Inc. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.lambda;

import java.io.Closeable;

import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.util.ShutdownHookManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.cloudera.oryx.common.lang.OryxShutdownHook;

/**
 * Hadoop-related utility methods.
 */
public final class HadoopUtils {

  private static final Logger log = LoggerFactory.getLogger(HadoopUtils.class);

  private static final OryxShutdownHook SHUTDOWN_HOOK = new OryxShutdownHook();

  private HadoopUtils() {}

  /**
   * Adds a shutdown hook that tries to call {@link Closeable#close()} on the given argument
   * at JVM shutdown. This integrates with Hadoop's {@link ShutdownHookManager} in order to
   * better interact with Spark's usage of the same.
   *
   * @param closeable thing to close
   */
  public static void closeAtShutdown(Closeable closeable) {
    if (SHUTDOWN_HOOK.addCloseable(closeable)) {
      try {
        // Spark uses SHUTDOWN_HOOK_PRIORITY + 30; this tries to execute earlier
        ShutdownHookManager.get().addShutdownHook(SHUTDOWN_HOOK, FileSystem.SHUTDOWN_HOOK_PRIORITY + 40);
      } catch (IllegalStateException ise) {
        log.warn("Can't close {} at shutdown since shutdown is in progress", closeable);
      }
    }
  }

}


/*
 * Copyright (c) 2016, Cloudera, Inc. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.api.speed;

import java.io.IOException;
import java.util.Iterator;
import java.util.Objects;

import org.apache.hadoop.conf.Configuration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.cloudera.oryx.api.KeyMessage;

/**
 * Convenience implementation of {@link SpeedModelManager} that provides default implementations.
 *
 * @param <K> type of key read from input topic
 * @param <M> type of message read from input topic
 * @param <U> type of update message read/written
 * @since 2.3.0
 */
public abstract class AbstractSpeedModelManager<K,M,U> implements SpeedModelManager<K,M,U> {

  private static final Logger log = LoggerFactory.getLogger(AbstractSpeedModelManager.class);

  @Override
  public void consume(Iterator<KeyMessage<String,U>> updateIterator, Configuration hadoopConf) throws IOException {
    while (updateIterator.hasNext()) {
      try {
        KeyMessage<String,U> km = updateIterator.next();
        String key = km.getKey();
        U message = km.getMessage();
        Objects.requireNonNull(key);
        consumeKeyMessage(key, message, hadoopConf);
      } catch (Exception e) {
        log.warn("Exception while processing message; continuing", e);
      }
    }
  }

  /**
   * Convenience method that is called by the default implementation of
   * {@link #consume(Iterator, Configuration)}, to process one key-message pair.
   * It does nothing, except log the message. This should generally be overridden
   * if and only if {@link #consume(Iterator, Configuration)} is not.
   *
   * @param key key to process (non-null)
   * @param message message to process
   * @param hadoopConf Hadoop configuration for process
   * @throws IOException if an error occurs while processing the message
   * @since 2.3.0
   */
  public void consumeKeyMessage(String key, U message, Configuration hadoopConf) throws IOException {
    log.info("{} : {}", key, message);
  }

}


/*
 * Copyright (c) 2014, Cloudera, Inc. and Intel Corp. All Rights Reserved.
 *
 * Cloudera, Inc. licenses this file to you under the Apache License,
 * Version 2.0 (the "License"). You may not use this file except in
 * compliance with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied. See the License for
 * the specific language governing permissions and limitations under the
 * License.
 */

package com.cloudera.oryx.app.serving.als.model;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Predicate;
import java.util.stream.Stream;

import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.eclipse.collections.api.map.MutableMap;
import org.eclipse.collections.api.set.MutableSet;
import org.eclipse.collections.impl.map.mutable.UnifiedMap;
import org.eclipse.collections.impl.map.mutable.primitive.ObjectIntHashMap;
import org.eclipse.collections.impl.set.mutable.UnifiedSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.cloudera.oryx.api.serving.ServingModel;
import com.cloudera.oryx.app.als.FeatureVectorsPartition;
import com.cloudera.oryx.app.als.PartitionedFeatureVectors;
import com.cloudera.oryx.app.als.RescorerProvider;
import com.cloudera.oryx.app.als.SolverCache;
import com.cloudera.oryx.app.serving.als.CosineDistanceSensitiveFunction;
import com.cloudera.oryx.common.collection.Pair;
import com.cloudera.oryx.common.collection.Pairs;
import com.cloudera.oryx.common.lang.AutoLock;
import com.cloudera.oryx.common.lang.AutoReadWriteLock;
import com.cloudera.oryx.common.lang.ToDoubleObjDoubleBiFunction;
import com.cloudera.oryx.common.math.Solver;

/**
 * Contains all data structures needed to serve real-time requests for an ALS-based recommender.
 */
public final class ALSServingModel implements ServingModel {

  private static final Logger log = LoggerFactory.getLogger(ALSServingModel.class);

  private static final ExecutorService executor = Executors.newFixedThreadPool(
      Runtime.getRuntime().availableProcessors(),
      new ThreadFactoryBuilder().setDaemon(true).setNameFormat("ALSServingModel-%d").build());

  private final LocalitySensitiveHash lsh;
  /** User-feature matrix. */
  private final FeatureVectorsPartition X;
  /** Item-feature matrix. This is partitioned into several maps for parallel access. */
  private final PartitionedFeatureVectors Y;
  /** Remembers items that each user has interacted with*/
  private final MutableMap<String,MutableSet<String>> knownItems; // Right now no corresponding "knownUsers" object
  private final AutoReadWriteLock knownItemsLock;
  private final MutableSet<String> expectedUserIDs;
  private final AutoReadWriteLock expectedUserIDsLock;
  private final MutableSet<String> expectedItemIDs;
  private final AutoReadWriteLock expectedItemIDsLock;
  private final SolverCache cachedYTYSolver;
  /** Number of features used in the model. */
  private final int features;
  /** Whether model uses implicit feedback. */
  private final boolean implicit;
  private final RescorerProvider rescorerProvider;

  /**
   * Creates an empty model.
   *
   * @param features number of features expected for user/item feature vectors
   * @param implicit whether model implements implicit feedback
   * @param sampleRate consider only approximately this fraction of all items when making recommendations.
   *  Candidates are chosen intelligently with locality sensitive hashing.
   * @param rescorerProvider optional instance of a {@link RescorerProvider}
   */
  ALSServingModel(int features, boolean implicit, double sampleRate, RescorerProvider rescorerProvider) {
    Preconditions.checkArgument(features > 0);
    Preconditions.checkArgument(sampleRate > 0.0 && sampleRate <= 1.0);

    lsh = new LocalitySensitiveHash(sampleRate, features);

    X = new FeatureVectorsPartition();
    Y = new PartitionedFeatureVectors(
        lsh.getNumPartitions(),
        executor,
        (String id, float[] vector) -> lsh.getIndexFor(vector));

    knownItems = UnifiedMap.newMap();
    knownItemsLock = new AutoReadWriteLock();

    expectedUserIDs = UnifiedSet.newSet();
    expectedUserIDsLock = new AutoReadWriteLock();
    expectedItemIDs = UnifiedSet.newSet();
    expectedItemIDsLock = new AutoReadWriteLock();

    cachedYTYSolver = new SolverCache(executor, Y);

    this.features = features;
    this.implicit = implicit;
    this.rescorerProvider = rescorerProvider;
  }

  public int getFeatures() {
    return features;
  }

  public boolean isImplicit() {
    return implicit;
  }

  public RescorerProvider getRescorerProvider() {
    return rescorerProvider;
  }

  public float[] getUserVector(String user) {
    return X.getVector(user);
  }

  public float[] getItemVector(String item) {
    return Y.getVector(item);
  }

  void setUserVector(String user, float[] vector) {
    Preconditions.checkArgument(vector.length == features);
    X.setVector(user, vector);
    try (AutoLock al = expectedUserIDsLock.autoWriteLock()) {
      expectedUserIDs.remove(user);
    }
  }

  void setItemVector(String item, float[] vector) {
    Preconditions.checkArgument(vector.length == features);
    Y.setVector(item, vector);
    try (AutoLock al = expectedItemIDsLock.autoWriteLock()) {
      expectedItemIDs.remove(item);
    }
    // Not clear if it's too inefficient to clear and recompute YtY solver every time any bit
    // of Y changes, but it's the most correct
    cachedYTYSolver.setDirty();
  }

  /**
   * @param user user to get known items for
   * @return set of known items for the user (immutable, but thread-safe)
   */
  public Set<String> getKnownItems(String user) {
    MutableSet<String> knownItemsForUser = doGetKnownItems(user);
    if (knownItemsForUser == null) {
      return Collections.emptySet();
    }
    synchronized (knownItemsForUser) {
      if (knownItemsForUser.isEmpty()) {
        return Collections.emptySet();
      }
      // Must copy since the original object is synchronized
      return knownItemsForUser.clone().asUnmodifiable();
    }
  }

  private MutableSet<String> doGetKnownItems(String user) {
    try (AutoLock al = knownItemsLock.autoReadLock()) {
      return knownItems.get(user);
    }
  }

  /**
   * @return mapping of user IDs to count of items the user has interacted with
   */
  public Map<String,Integer> getUserCounts() {
    Map<String,Integer> counts;
    try (AutoLock al = knownItemsLock.autoReadLock()) {
      counts = new HashMap<>(knownItems.size());
      knownItems.forEach((userID, ids) -> {
        int numItems;
        synchronized (ids) {
          numItems = ids.size();
        }
        counts.put(userID, numItems);
      });
    }
    return counts;
  }

  /**
   * @return mapping of item IDs to count of users that have interacted with that item
   */
  public Map<String,Integer> getItemCounts() {
    ObjectIntHashMap<String> counts = ObjectIntHashMap.newMap();
    try (AutoLock al = knownItemsLock.autoReadLock()) {
      knownItems.values().forEach(ids -> {
        synchronized (ids) {
          ids.forEach(id -> counts.addToValue(id, 1));
        }
      });
    }
    // No way to get Java map from primitive map directly (?)
    Map<String,Integer> javaCounts = new HashMap<>(counts.size());
    counts.forEachKeyValue(javaCounts::put);
    return javaCounts;
  }

  void addKnownItems(String user, Collection<String> items) {
    if (!items.isEmpty()) {
      MutableSet<String> knownItemsForUser = doGetKnownItems(user);

      if (knownItemsForUser == null) {
        try (AutoLock al = knownItemsLock.autoWriteLock()) {
          // Check again
          knownItemsForUser = knownItems.computeIfAbsent(user, k -> UnifiedSet.newSet());
        }
      }

      synchronized (knownItemsForUser) {
        knownItemsForUser.addAll(items);
      }
    }
  }

  /**
   * @param user user to get known item vectors for
   * @return {@code null} if the user is not known to the model, or if there are no known items for the user
   */
  public List<Pair<String,float[]>> getKnownItemVectorsForUser(String user) {
    float[] userVector = getUserVector(user);
    if (userVector == null) {
      return null;
    }
    Collection<String> knownItemsForUser = doGetKnownItems(user);
    if (knownItemsForUser == null) {
      return null;
    }
    synchronized (knownItemsForUser) {
      int size = knownItemsForUser.size();
      if (size == 0) {
        return null;
      }
      List<Pair<String,float[]>> idVectors = new ArrayList<>(size);
      for (String itemID : knownItemsForUser) {
        float[] vector = getItemVector(itemID);
        if (vector != null) {
          idVectors.add(new Pair<>(itemID, vector));
        }
      }
      return idVectors.isEmpty() ? null : idVectors;
    }
  }

  public Stream<Pair<String,Double>> topN(
      CosineDistanceSensitiveFunction scoreFn,
      ToDoubleObjDoubleBiFunction<String> rescoreFn,
      int howMany,
      Predicate<String> allowedPredicate) {
    int[] candidateIndices = lsh.getCandidateIndices(scoreFn.getTargetVector());
    Stream<Pair<String,Double>> stream = Y.mapPartitionsParallel(
        partition -> {
          TopNConsumer consumer = new TopNConsumer(howMany, scoreFn, rescoreFn, allowedPredicate);
          partition.forEach(consumer);
          return consumer.getTopN();
        },
        candidateIndices,
        false);
    return stream.sorted(Pairs.orderBySecond(Pairs.SortOrder.DESCENDING)).limit(howMany);
  }

  /**
   * @return all user IDs in the model
   */
  public Collection<String> getAllUserIDs() {
    Collection<String> allUserIDs = UnifiedSet.newSet(X.size());
    X.addAllIDsTo(allUserIDs);
    return allUserIDs;
  }

  /**
   * @return all item IDs in the model
   */
  public Collection<String> getAllItemIDs() {
    Collection<String> allItemIDs = UnifiedSet.newSet(Y.size());
    Y.addAllIDsTo(allItemIDs);
    return allItemIDs;
  }

  /**
   * @return a {@link Solver} for use in solving systems involving YT*Y
   */
  public Solver getYTYSolver() {
    return cachedYTYSolver.get(true);
  }

  void precomputeSolvers() {
    cachedYTYSolver.compute();
  }

  /**
   * Retains only users that are expected to appear
   * in the upcoming model updates, or, that have arrived recently. This also clears the
   * recent known users data structure.
   *
   * @param users users that should be retained, which are coming in the new model updates
   */
  void retainRecentAndUserIDs(Collection<String> users) {
    X.retainRecentAndIDs(users);
    try (AutoLock al = expectedUserIDsLock.autoWriteLock()) {
      expectedUserIDs.clear();
      expectedUserIDs.addAll(users);
      X.removeAllIDsFrom(expectedUserIDs);
    }
  }

  /**
   * Retains only items that are expected to appear
   * in the upcoming model updates, or, that have arrived recently. This also clears the
   * recent known items data structure.
   *
   * @param items items that should be retained, which are coming in the new model updates
   */
  void retainRecentAndItemIDs(Collection<String> items) {
    Y.retainRecentAndIDs(items);
    try (AutoLock al = expectedItemIDsLock.autoWriteLock()) {
      expectedItemIDs.clear();
      expectedItemIDs.addAll(items);
      Y.removeAllIDsFrom(expectedItemIDs);
    }
  }

  /**
   * Like {@link #retainRecentAndUserIDs(Collection)} and {@link #retainRecentAndItemIDs(Collection)}
   * but affects the known-items data structure.
   *
   * @param users users that should be retained, which are coming in the new model updates
   * @param items items that should be retained, which are coming in the new model updates
   */
  void retainRecentAndKnownItems(Collection<String> users, Collection<String> items) {
    // Keep all users in the new model, or, that have been added since last model
    MutableSet<String> recentUserIDs = UnifiedSet.newSet();
    X.addAllRecentTo(recentUserIDs);
    try (AutoLock al = knownItemsLock.autoWriteLock()) {
      knownItems.keySet().removeIf(key -> !users.contains(key) && !recentUserIDs.contains(key));
    }

    // This will be easier to quickly copy the whole (smallish) set rather than
    // deal with locks below
    MutableSet<String> allRecentKnownItems = UnifiedSet.newSet();
    Y.addAllRecentTo(allRecentKnownItems);

    Predicate<String> notKeptOrRecent = value -> !items.contains(value) && !allRecentKnownItems.contains(value);
    try (AutoLock al = knownItemsLock.autoReadLock()) {
      knownItems.values().forEach(knownItemsForUser -> {
        synchronized (knownItemsForUser) {
          knownItemsForUser.removeIf(notKeptOrRecent);
        }
      });
    }
  }

  /**
   * @return number of users in the model
   */
  public int getNumUsers() {
    return X.size();
  }

  /**
   * @return number of items in the model
   */
  public int getNumItems() {
    return Y.size();
  }

  @Override
  public float getFractionLoaded() {
    int expected = 0;
    try (AutoLock al = expectedUserIDsLock.autoReadLock()) {
      expected += expectedUserIDs.size();
    }
    try (AutoLock al = expectedItemIDsLock.autoReadLock()) {
      expected += expectedItemIDs.size();
    }
    if (expected == 0) {
      return 1.0f;
    }
    float loaded = (float) getNumUsers() + getNumItems();
    return loaded / (loaded + expected);
  }

  @Override
  public String toString() {
    return "ALSServingModel[features:" + features + ", implicit:" + implicit +
        ", X:(" + getNumUsers() + " users), Y:(" + getNumItems() + " items, partitions: " +
        Y + "...), fractionLoaded:" + getFractionLoaded() + "]";
  }

}