package org.ohdsi.webapi.service;

import static org.ohdsi.webapi.util.SecurityUtils.whitelist;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;

import java.util.stream.Collectors;
import javax.annotation.PostConstruct;
import javax.ws.rs.*;
import javax.ws.rs.core.MediaType;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.math.NumberUtils;
import org.ohdsi.circe.helper.ResourceHelper;
import org.ohdsi.webapi.cohortanalysis.CohortAnalysis;
import org.ohdsi.webapi.cohortanalysis.CohortAnalysisTask;
import org.ohdsi.webapi.cohortanalysis.CohortSummary;
import org.ohdsi.webapi.cohortdefinition.CohortDefinitionRepository;
import org.ohdsi.webapi.cohortdefinition.dto.CohortDTO;
import org.ohdsi.webapi.cohortresults.*;
import org.ohdsi.webapi.cohortresults.mapper.AnalysisResultsMapper;
import org.ohdsi.webapi.model.results.Analysis;
import org.ohdsi.webapi.model.results.AnalysisResults;
import org.ohdsi.webapi.source.Source;
import org.ohdsi.webapi.source.SourceDaimon;
import org.ohdsi.webapi.util.PreparedStatementRenderer;
import org.ohdsi.webapi.util.SessionUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.ResultSetExtractor;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.stereotype.Component;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.ByteArrayOutputStream;
import java.sql.ResultSetMetaData;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import javax.ws.rs.core.Response;

import org.ohdsi.webapi.person.CohortPerson;

/**
 * REST Services related to retrieving
 * cohort analysis (a.k.a Heracles Results) analyses results. 
 * More information on the Heracles project
 * can be found at {@link https://www.ohdsi.org/web/wiki/doku.php?id=documentation:software:heracles}.
 * The implementation found in WebAPI represents a migration of the functionality
 * from the stand-alone HERACLES application to integrate it into WebAPI and
 * ATLAS.
 * 
 * @summary Cohort Analysis Results (a.k.a Heracles Results)
 */
@Path("/cohortresults")
@Component
public class CohortResultsService extends AbstractDaoService {

  public static final String MIN_COVARIATE_PERSON_COUNT = "10";
  public static final String MIN_INTERVAL_PERSON_COUNT = "10";

  public static final String BASE_SQL_PATH = "/resources/cohortresults/sql";

  @Autowired
  private VisualizationDataRepository visualizationDataRepository;

  @Autowired
  private CohortDefinitionService cohortDefinitionService;

  @Autowired
  private CohortDefinitionRepository cohortDefinitionRepository;

  @Autowired
  private ObjectMapper mapper;

  private CohortResultsAnalysisRunner queryRunner = null;

  @PostConstruct
  public void init() {
    queryRunner = new CohortResultsAnalysisRunner(this.getSourceDialect(), this.visualizationDataRepository, mapper);
  }

  /**
   * Queries for cohort analysis results for the given cohort definition id
   *
   * @summary Get results for analysis group
   * @param id cohort_defintion id
   * @param analysisGroup Name of the analysisGrouping under the
   * /resources/cohortresults/sql/ directory
   * @param analysisName Name of the analysis, currently the same name as the
   * sql file under analysisGroup
   * @param sourceKey the source to retrieve results
   * @return List of key, value pairs
   */
  @GET
  @Path("{sourceKey}/{id}/raw/{analysis_group}/{analysis_name}")
  @Produces(MediaType.APPLICATION_JSON)
  public List<Map<String, String>> getCohortResultsRaw(@PathParam("id") final int id, @PathParam("analysis_group") final String analysisGroup,
          @PathParam("analysis_name") final String analysisName,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @PathParam("sourceKey") String sourceKey) {
    List<Map<String, String>> results;
    String sqlPath = BASE_SQL_PATH + "/" + analysisGroup + "/" + analysisName + ".sql";

    Source source = getSourceRepository().findBySourceKey(sourceKey);
    try {
      PreparedStatementRenderer psr = prepareGetCohortResultsRaw(id, minCovariatePersonCountParam,
        minIntervalPersonCountParam, sqlPath, source);
      return genericResultSetLoader(psr, source);
    } catch (Exception e) {
      log.error("Unable to translate sql for analysis {}", analysisName, e);
      return null;
    }
  }

  protected PreparedStatementRenderer prepareGetCohortResultsRaw(final int id,
                                                                 final Integer minCovariatePersonCountParam,
                                                                 final Integer minIntervalPersonCountParam, String sqlPath,
                                                                 Source source) {

    String resourcePath = sqlPath;
    String vocabularyTableQualifier = source.getTableQualifier(SourceDaimon.DaimonType.Vocabulary);
    String resultsTableQualifier = source.getTableQualifier(SourceDaimon.DaimonType.Results);


    String[] searchStringNames = new String[]{"cdm_database_schema", "ohdsi_database_schema"};
    String[] replacementNames = new String[]{vocabularyTableQualifier, resultsTableQualifier};


    String[] variableNames = new String[]{"cohortDefinitionId", "minCovariatePersonCount", "minIntervalPersonCount"};
    Object[] variableValues = new Object[]
        {id, (minCovariatePersonCountParam == null ? MIN_COVARIATE_PERSON_COUNT : minCovariatePersonCountParam),
            (minIntervalPersonCountParam == null ? MIN_INTERVAL_PERSON_COUNT : minIntervalPersonCountParam)};

    return new PreparedStatementRenderer(source, resourcePath, searchStringNames, replacementNames, variableNames, variableValues);
  }

  /**
   * Export the cohort analysis results to a ZIP file
   * 
   * @summary Export cohort analysis results
   * @param id The cohort ID
   * @param sourceKey The source Key
   * @return A response containing the .ZIP file of results
   */
  @GET
  @Path("{sourceKey}/{id}/export.zip")
  @Produces(MediaType.APPLICATION_OCTET_STREAM)
  public Response exportCohortResults(@PathParam("id") int id, @PathParam("sourceKey") String sourceKey) {
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    ZipOutputStream zos = new ZipOutputStream(baos);

    try {
      Source source = getSourceRepository().findBySourceKey(sourceKey);
      String sqlPath = BASE_SQL_PATH + "/raw/getAllResults.sql";
      String tqName = "tableQualifier";
      String tqValue = source.getTableQualifier(SourceDaimon.DaimonType.Results);
      PreparedStatementRenderer psr = new PreparedStatementRenderer(source, sqlPath, tqName, tqValue, "cohortDefinitionId", whitelist(id));

      final StringBuilder resultData = new StringBuilder();
      final StringBuilder resultDistributionData = new StringBuilder();

      // results export
      getSourceJdbcTemplate(source).query(psr.getSql(), psr.getSetter(), new RowMapper<Void>() {
        @Override
        public Void mapRow(ResultSet rs, int arg1) throws SQLException {

          ResultSetMetaData metaData = rs.getMetaData();
          int colCount = metaData.getColumnCount();
          for (int i = 1; i <= colCount; i++) {
            if (i > 1) {
              resultData.append("\t");
            }
            resultData.append(String.valueOf(rs.getObject(i)));
          }
          resultData.append("\r\n");
          return null;
        }
      });

      ZipEntry resultsEntry = new ZipEntry("cohort_" + String.valueOf(id) + "_results.tsv");

      zos.putNextEntry(resultsEntry);
      zos.write(resultData.toString().getBytes());
      zos.closeEntry();

      // result distribution export
      sqlPath = BASE_SQL_PATH + "/raw/getAllResultDistributions.sql";
      psr = new PreparedStatementRenderer(source, sqlPath, tqName, tqValue, "cohortDefinitionId", whitelist(id));
      getSourceJdbcTemplate(source).query(psr.getSql(), psr.getSetter(), new RowMapper<Void>() {
        @Override
        public Void mapRow(ResultSet rs, int arg1) throws SQLException {
          ResultSetMetaData metaData = rs.getMetaData();
          int colCount = metaData.getColumnCount();
          for (int i = 1; i <= colCount; i++) {
            if (i > 1) {
              resultDistributionData.append("\t");
            }
            resultDistributionData.append(String.valueOf(rs.getObject(i)));
          }
          resultDistributionData.append("\r\n");
          return null;
        }
      });

      ZipEntry resultsDistEntry = new ZipEntry("cohort_" + String.valueOf(id) + "_results_dist.tsv");
      zos.putNextEntry(resultsDistEntry);
      zos.write(resultDistributionData.toString().getBytes());
      zos.closeEntry();

      // include cohort definition in export
      CohortDTO cohortDefinition = cohortDefinitionService.getCohortDefinition(id);
      ByteArrayOutputStream cohortDefinitionStream = new ByteArrayOutputStream();
      mapper.writeValue(cohortDefinitionStream, cohortDefinition);
      cohortDefinitionStream.flush();

      ZipEntry cohortDefinitionEntry = new ZipEntry("cohort_" + String.valueOf(id) + "_definition.json");
      zos.putNextEntry(cohortDefinitionEntry);
      zos.write(cohortDefinitionStream.toByteArray());
      zos.closeEntry();

      zos.close();
      baos.flush();
      baos.close();
    } catch (Exception ex) {
      throw new RuntimeException(ex);
    }

    Response response = Response
            .ok(baos)
            .type(MediaType.APPLICATION_OCTET_STREAM)
            .build();

    return response;
  }

  /**
   * Provides a warmup mechanism for the data visualization cache. This
   * endpoint does not appear to be used and may be a hold over from the 
   * original HERACLES implementation
   * 
   * @summary Warmup data visualizations
   * @param task The cohort analysis task
   * @return The number of report visualizations warmed
   */
  @POST
  @Path("/warmup")
  @Produces(MediaType.APPLICATION_JSON)
  @Consumes(MediaType.APPLICATION_JSON)
  public int warmUpVisualizationData(CohortAnalysisTask task) {
    return this.queryRunner.warmupData(this.getSourceJdbcTemplate(task.getSource()), task);

  }

  /**
   * Provides a list of cohort analysis visualizations that are completed
   * 
   * @summary Get completed cohort analysis visualizations
   * @param id The cohort ID
   * @param sourceKey The source key
   * @return A list of visualization keys that are complete
   */
  @GET
  @Path("{sourceKey}/{id}/completed")
  @Produces(MediaType.APPLICATION_JSON)
  @Consumes(MediaType.APPLICATION_JSON)
  public Collection<String> getCompletedVisualiztion(@PathParam("id") final int id,
          @PathParam("sourceKey") final String sourceKey) {
    Source source = getSourceRepository().findBySourceKey(sourceKey);
    List<VisualizationData> vizData = this.visualizationDataRepository.findByCohortDefinitionIdAndSourceId(id, source.getSourceId());
    Set<String> completed = new HashSet<>();
    if (CollectionUtils.isNotEmpty(vizData)) {
      for (VisualizationData viz : vizData) {
        completed.add(viz.getVisualizationKey());
      }
    }
    return completed;
  }


  /**
   * Retrieves the tornado plot
   * 
   * @summary Get the tornado plot
   * @param sourceKey The source key
   * @param cohortDefinitionId The cohort definition id
   * @return The tornado plot data
   */
  @GET
  @Path("{sourceKey}/{id}/tornado")
  @Produces(MediaType.APPLICATION_JSON)
  public TornadoReport getTornadoReport(@PathParam("sourceKey") final String sourceKey, @PathParam("id") final int cohortDefinitionId) {
        Source source = getSourceRepository().findBySourceKey(sourceKey);
        TornadoReport tornadoReport = new TornadoReport();
        tornadoReport.tornadoRecords = queryRunner.getTornadoRecords(getSourceJdbcTemplate(source), cohortDefinitionId, source);
        tornadoReport.profileSamples = queryRunner.getProfileSampleRecords(getSourceJdbcTemplate(source), cohortDefinitionId, source);
        return tornadoReport;
  }

  /**
   * Queries for cohort analysis dashboard for the given cohort definition id
   *
   * @summary Get the dashboard
   * @param id The cohort definition id
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param demographicsOnly only render gender and age
   * @return CohortDashboard
   */
  @GET
  @Path("{sourceKey}/{id}/dashboard")
  @Produces(MediaType.APPLICATION_JSON)
  public CohortDashboard getDashboard(@PathParam("id") final int id,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @QueryParam("demographics_only") final boolean demographicsOnly,
          @PathParam("sourceKey") final String sourceKey,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {

    final String key = CohortResultsAnalysisRunner.DASHBOARD;
    Source source = getSourceRepository().findBySourceKey(sourceKey);
    VisualizationData data = refresh ? null : this.visualizationDataRepository.findByCohortDefinitionIdAndSourceIdAndVisualizationKey(id, source.getSourceId(), key);

    CohortDashboard dashboard = null;

    if (refresh || data == null) {
      dashboard = queryRunner.getDashboard(getSourceJdbcTemplate(source), id, source,
              minCovariatePersonCountParam, minIntervalPersonCountParam, demographicsOnly, true);

    } else {
      try {
        dashboard = mapper.readValue(data.getData(), CohortDashboard.class);
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return dashboard;

  }

  /**
   * Queries for cohort analysis condition treemap results for the given cohort
   * definition id
   * 
   * @summary Get condition treemap
   * @param sourceKey The source key
   * @param id The cohort ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param refresh Boolean - refresh visualization data
   * @return List<HierarchicalConceptRecord>
   */
  @GET
  @Path("{sourceKey}/{id}/condition/")
  @Produces(MediaType.APPLICATION_JSON)
  public List<HierarchicalConceptRecord> getConditionTreemap(@PathParam("sourceKey") String sourceKey, @PathParam("id") final int id,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {

    Source source = getSourceRepository().findBySourceKey(sourceKey);
    final String key = CohortResultsAnalysisRunner.CONDITION;
    List<HierarchicalConceptRecord> res = null;
    VisualizationData data = refresh ? null : this.visualizationDataRepository.findByCohortDefinitionIdAndSourceIdAndVisualizationKey(id, source.getSourceId(), key);

    if (refresh || data == null) {
      res = this.queryRunner.getConditionTreemap(this.getSourceJdbcTemplate(source), id, minCovariatePersonCountParam, minIntervalPersonCountParam, source, true);
    } else {
      try {
        res = mapper.readValue(data.getData(), new TypeReference<List<HierarchicalConceptRecord>>() {
        });
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return res;
  }

  /**
   * Get the distinct person count for a cohort
   * 
   * @summary Get distinct person count
   * @param sourceKey The source key
   * @param id The cohort ID
   * @param refresh Boolean - refresh visualization data
   * @return Distinct person count as integer
   */
  @GET
  @Path("{sourceKey}/{id}/distinctPersonCount/")
  @Produces(MediaType.APPLICATION_JSON)
  public Integer getRawDistinctPersonCount(@PathParam("sourceKey") String sourceKey,
          @PathParam("id") String id,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {
    Source source = getSourceRepository().findBySourceKey(sourceKey);
    PreparedStatementRenderer psr = prepareGetRawDistinctPersonCount(id, source);
    Integer result = getSourceJdbcTemplate(source).query(psr.getSql(), psr.getSetter(), new ResultSetExtractor<Integer>() {
      @Override
      public Integer extractData(ResultSet rs) throws SQLException {

        while (rs.next()) {
          return rs.getInt(1);
        }
        return null;
      }
    });
    return result;
  }

  protected PreparedStatementRenderer prepareGetRawDistinctPersonCount(String id, Source source) {

    String sqlPath = BASE_SQL_PATH + "/raw/getTotalDistinctPeople.sql";
    String tqValue = source.getTableQualifier(SourceDaimon.DaimonType.Results);
    String tqName = "tableQualifier";
    return new PreparedStatementRenderer(source, sqlPath, tqName, tqValue, "id", Integer.valueOf(id));
  }

  /**
   * Queries for cohort analysis condition drilldown results for the given
   * cohort definition id and condition id
   * 
   * @param sourceKey The source key
   * @param id The cohort ID
   * @param conditionId The condition concept ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param refresh Boolean - refresh visualization data
   * @return The CohortConditionDrilldown detail object
   */
  @GET
  @Path("{sourceKey}/{id}/condition/{conditionId}")
  @Produces(MediaType.APPLICATION_JSON)
  public CohortConditionDrilldown getConditionResults(@PathParam("sourceKey") String sourceKey,
          @PathParam("id") final int id,
          @PathParam("conditionId") final int conditionId,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {
    CohortConditionDrilldown drilldown = null;
    final String key = CohortResultsAnalysisRunner.CONDITION_DRILLDOWN;
    Source source = getSourceRepository().findBySourceKey(sourceKey);
    VisualizationData data = refresh ? null : this.visualizationDataRepository.findByCohortDefinitionIdAndSourceIdAndVisualizationKeyAndDrilldownId(id, source.getSourceId(), key, conditionId);

    if (refresh || data == null) {
      drilldown = this.queryRunner.getConditionResults(this.getSourceJdbcTemplate(source), id, conditionId, minCovariatePersonCountParam, minIntervalPersonCountParam, source, true);
    } else {
      try {
        drilldown = mapper.readValue(data.getData(), CohortConditionDrilldown.class);
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return drilldown;

  }

  /**
   * Queries for cohort analysis condition era treemap results for the given
   * cohort definition id
   * 
   * @param sourceKey The source key
   * @param id The cohort ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param refresh Boolean - refresh visualization data
   * @return List<HierarchicalConceptRecord>
   */
  @GET
  @Path("{sourceKey}/{id}/conditionera/")
  @Produces(MediaType.APPLICATION_JSON)
  public List<HierarchicalConceptRecord> getConditionEraTreemap(@PathParam("sourceKey") final String sourceKey,
          @PathParam("id") final int id,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {

    Source source = getSourceRepository().findBySourceKey(sourceKey);
    final String key = CohortResultsAnalysisRunner.CONDITION_ERA;
    VisualizationData data = refresh ? null : this.visualizationDataRepository.findByCohortDefinitionIdAndSourceIdAndVisualizationKey(id, source.getSourceId(), key);

    List<HierarchicalConceptRecord> res = null;

    if (data == null || refresh) {
      res = this.queryRunner.getConditionEraTreemap(this.getSourceJdbcTemplate(source), id, minCovariatePersonCountParam, minIntervalPersonCountParam, source, true);
    } else {
      try {
        res = mapper.readValue(data.getData(), new TypeReference<List<HierarchicalConceptRecord>>() {
        });
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return res;
  }

  /**
   * Get the completed analyses IDs for the selected cohort and source key
   * 
   * @summary Get completed analyses IDs
   * @param sourceKey The source key
   * @param id The cohort ID
   * @return A list of completed analysis IDs
   */
  @GET
  @Path("{sourceKey}/{id}/analyses")
  @Produces(MediaType.APPLICATION_JSON)
  public List<Integer> getCompletedAnalyses(@PathParam("sourceKey") String sourceKey, @PathParam("id") String id) {
    Source source = getSourceRepository().findBySourceKey(sourceKey);
		int sourceId = source.getSourceId();
		
    PreparedStatementRenderer psr = prepareGetCompletedAnalysis(id, sourceId);
    final String sql = psr.getSql();
    return this.getJdbcTemplate().query(sql, psr.getSetter(), new RowMapper<Integer>() {
          @Override
          public Integer mapRow(ResultSet resultSet, int arg1) throws SQLException {

            return resultSet.getInt(1);
          }
        }
    );
  }

  class GenerationInfoDTO {
    private String sourceKey;
    private Integer analysisId;
    private Integer progress;

    public GenerationInfoDTO() {
    }

    public GenerationInfoDTO(String sourceKey, Integer analysisId, Integer progress) {
      this.sourceKey = sourceKey;
      this.analysisId = analysisId;
      this.progress = progress;
    }

    public String getSourceKey() {
      return sourceKey;
    }

    public void setSourceKey(String sourceKey) {
      this.sourceKey = sourceKey;
    }

    public Integer getAnalysisId() {
      return analysisId;
    }

    public void setAnalysisId(Integer analysisId) {
      this.analysisId = analysisId;
    }

    public Integer getProgress() {
      return progress;
    }

    public void setProgress(Integer progress) {
      this.progress = progress;
    }
  }

  /**
   * Get the analysis generation progress
   * 
   * @summary Get analysis progress
   * @param sourceKey The source key
   * @param id The cohort ID
   * @return The generation progress information
   */
  @GET
  @Path("{sourceKey}/{id}/info")
  @Produces(MediaType.APPLICATION_JSON)
  public GenerationInfoDTO getAnalysisProgress(@PathParam("sourceKey") String sourceKey, @PathParam("id") Integer id) {

    return getTransactionTemplateRequiresNew().execute(status -> {
      org.ohdsi.webapi.cohortdefinition.CohortDefinition def = cohortDefinitionRepository.findOne(id);
      Source source = getSourceRepository().findBySourceKey(sourceKey);
      return def.getCohortAnalysisGenerationInfoList().stream()
              .filter(cd -> Objects.equals(cd.getSourceId(), source.getSourceId()))
              .findFirst().map(gen -> new GenerationInfoDTO(sourceKey, id, gen.getProgress()))
              .<RuntimeException>orElseThrow(NotFoundException::new);
    });
  }

  protected PreparedStatementRenderer prepareGetCompletedAnalysis(String id, int sourceId) {

    String sqlPath = BASE_SQL_PATH + "/raw/getCompletedAnalyses.sql";
    PreparedStatementRenderer psr = new PreparedStatementRenderer(getSourceRepository().findBySourceId(sourceId)
			, sqlPath
			, new String[]{"tableQualifier"}, new String[] { this.getOhdsiSchema()}
			, new String[]{"cohort_definition_id", "source_id"}, new Object[]{Integer.valueOf(id), Integer.valueOf(sourceId)});
    return psr;
  }

  /**
   * Queries for cohort analysis condition era drilldown results for the given
   * cohort definition id and condition id
   * 
   * @summary Get condition era drilldown report
   * @param id The cohort ID
   * @param conditionId The condition ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param sourceKey The source key
   * @param refresh Boolean - refresh visualization data
   * @return The CohortConditionEraDrilldown object
   */
  @GET
  @Path("{sourceKey}/{id}/conditionera/{conditionId}")
  @Produces(MediaType.APPLICATION_JSON)
  public CohortConditionEraDrilldown getConditionEraDrilldown(@PathParam("id") final int id,
                                                              @PathParam("conditionId") final int conditionId,
                                                              @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
                                                              @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
                                                              @PathParam("sourceKey") final String sourceKey,
                                                              @DefaultValue("false") @QueryParam("refresh") boolean refresh) {

    CohortConditionEraDrilldown drilldown = null;
    final String key = CohortResultsAnalysisRunner.CONDITION_ERA_DRILLDOWN;
    Source source = getSourceRepository().findBySourceKey(sourceKey);
    VisualizationData data = refresh ? null : this.visualizationDataRepository
            .findByCohortDefinitionIdAndSourceIdAndVisualizationKeyAndDrilldownId(id, source.getSourceId(), key, conditionId);

    if (refresh || data == null) {
      drilldown = this.queryRunner.getConditionEraDrilldown(this.getSourceJdbcTemplate(source), id, conditionId, minCovariatePersonCountParam, minIntervalPersonCountParam, source, true);
    } else {
      try {
        drilldown = mapper.readValue(data.getData(), CohortConditionEraDrilldown.class);
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return drilldown;

  }

  /**
   * Queries for drug analysis treemap results for the given cohort
   * definition id
   * 
   * @summary Get drug treemap
   * @param id The cohort ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param sourceKey The source key
   * @param refresh Boolean - refresh visualization data
   * @return List<HierarchicalConceptRecord>
   */
  @GET
  @Path("{sourceKey}/{id}/drug/")
  @Produces(MediaType.APPLICATION_JSON)
  public List<HierarchicalConceptRecord> getDrugTreemap(@PathParam("id") final int id,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @PathParam("sourceKey") final String sourceKey,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {

    Source source = getSourceRepository().findBySourceKey(sourceKey);
    final String key = CohortResultsAnalysisRunner.DRUG;
    VisualizationData data = refresh ? null : this.visualizationDataRepository.findByCohortDefinitionIdAndSourceIdAndVisualizationKey(id, source.getSourceId(), key);

    List<HierarchicalConceptRecord> res = null;
    if (refresh || data == null) {
      res = this.queryRunner.getDrugTreemap(this.getSourceJdbcTemplate(source), id, minCovariatePersonCountParam, minIntervalPersonCountParam, source, true);
    } else {
      try {
        res = mapper.readValue(data.getData(), new TypeReference<List<HierarchicalConceptRecord>>() {
        });
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return res;
  }

  /**
   *
   * @param id cohort_defintion id
   * @param drugId drug_id (from concept)
   * @return CohortDrugDrilldown
   */
  /**
   * Queries for cohort analysis drug drilldown results for the given cohort
   * definition id and drug id
   * 
   * @summary Get drug drilldown report
   * @param id The cohort ID
   * @param drugId The drug concept ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param sourceKey The source key
   * @param refresh Boolean - refresh visualization data
   * @return 
   */
  @GET
  @Path("{sourceKey}/{id}/drug/{drugId}")
  @Produces(MediaType.APPLICATION_JSON)
  public CohortDrugDrilldown getDrugResults(@PathParam("id") final int id, @PathParam("drugId") final int drugId,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @PathParam("sourceKey") final String sourceKey,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {
    CohortDrugDrilldown drilldown = null;
    final String key = CohortResultsAnalysisRunner.DRUG_DRILLDOWN;
    Source source = getSourceRepository().findBySourceKey(sourceKey);
    VisualizationData data = refresh ? null : this.visualizationDataRepository.findByCohortDefinitionIdAndSourceIdAndVisualizationKeyAndDrilldownId(id, source.getSourceId(), key, drugId);

    if (refresh || data == null) {
      drilldown = this.queryRunner.getDrugResults(this.getSourceJdbcTemplate(source), id, drugId, minCovariatePersonCountParam, minIntervalPersonCountParam, source, true);
    } else {
      try {
        drilldown = mapper.readValue(data.getData(), CohortDrugDrilldown.class);
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return drilldown;

  }

  /**
   * Queries for cohort analysis drug era treemap results for the given cohort
   * definition id
   * 
   * @summary Get drug era treemap report
   * @param id The cohort ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param sourceKey The source key
   * @param refresh Boolean - refresh visualization data
   * @return List<HierarchicalConceptRecord>
   */
  @GET
  @Path("{sourceKey}/{id}/drugera/")
  @Produces(MediaType.APPLICATION_JSON)
  public List<HierarchicalConceptRecord> getDrugEraTreemap(@PathParam("id") final int id,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
          @QueryParam("min_interval_person_count") final Integer minIntervalPersonCountParam,
          @PathParam("sourceKey") final String sourceKey,
          @DefaultValue("false") @QueryParam("refresh") boolean refresh) {

    Source source = getSourceRepository().findBySourceKey(sourceKey);
    List<HierarchicalConceptRecord> res = null;
    final String key = CohortResultsAnalysisRunner.DRUG_ERA;
    VisualizationData data = refresh ? null : this.visualizationDataRepository.findByCohortDefinitionIdAndSourceIdAndVisualizationKey(id, source.getSourceId(), key);

    if (refresh || data == null) {
      res = this.queryRunner.getDrugEraTreemap(this.getSourceJdbcTemplate(source), id, minCovariatePersonCountParam, minIntervalPersonCountParam, source, true);
    } else {
      try {
        res = mapper.readValue(data.getData(), new TypeReference<List<HierarchicalConceptRecord>>() {
        });
      } catch (Exception e) {
        log.error(whitelist(e));
      }
    }

    return res;
  }

  /**
   *
   * @param id cohort_defintion id
   * @param drugId drug_id (from concept)
   * @return CohortDrugEraDrilldown
   */
  /**
   * Queries for cohort analysis drug era drilldown results for the given cohort
   * definition id and drug id
   * 
   * @summary Get drug era drilldown report
   * @param id The cohort ID
   * @param drugId The drug concept ID
   * @param minCovariatePersonCountParam The minimum number of covariates per person
   * @param minIntervalPersonCountParam The minimum interval person count
   * @param sourceKey The source key
   * @param refresh Boolean - refresh visualization data
   * @return CohortDrugEraDrilldown
   */
  @GET
  @Path("{sourceKey}/{id}/drugera/{drugId}")
  @Produces(MediaType.APPLICATION_JSON)
  public CohortDrugEraDrilldown getDrugEraResults(@PathParam("id") final int id, @PathParam("drugId") final int drugId,
          @QueryParam("min_covariate_person_count") final Integer minCovariatePersonCountParam,
